Method of learning under constraint by converting parameters to polar coordinates.

polar.py

import numpy as np

class PolarDiff:

    def __init__(self, n, r2=1, x=None, clip=None, epsilon=1e-7):

        if x is not None:
            if n != len(x):
                raise ValueError
            self.theta = np.zeros(n-1)
            tmp = r2
            for i, xi in enumerate(x[:-1]):
                self.theta[i] = np.arcsin(np.sqrt(xi/tmp))
                tmp *= np.cos(self.theta[i])**2
        else:
            self.theta = (np.random.random(size=n-1) - 0.5) * 2 * np.pi
        self.r2 = r2
        self.clip = clip
        self.epsilon = epsilon
        self.x = np.array(x)
        self.sin = np.zeros_like(self.theta)
        self.cos = np.zeros_like(self.theta)
        self.tan = np.zeros_like(self.theta)
        self.calc_sin_cos_tan()

    def calc_sin_cos_tan(self):

        self.sin = np.sin(self.theta)
        self.cos = np.cos(self.theta)
        self.tan = np.tan(self.theta)

    def calc_x(self):

        tmp = self.r2
        for i, theta in enumerate(self.theta):
            self.x[i] = tmp * np.sin(theta)**2
            tmp *= np.cos(theta)**2
        self.x[-1] = tmp
    
    def update(self, df, max_degree=10):

        n = len(self.x)
        pi2 = np.pi / 2
        grads = np.zeros_like(self.theta)
        for i in range(n - 1):
            grad = 0
            for j in range(i, n):
                if i == j:
                    grad += 2 * df[j] * self.x[j] / self.tan[i]
                else:
                    grad -= 2 * df[j] * self.x[j] * self.tan[i]
            
            if self.clip is not None:
                if self.clip < grad:
                    grad = self.clip
            grads[i] = grad
        lr = max_degree / np.max(np.abs(grads)) * np.pi / 180
        self.theta = self.theta - lr * grads
        self.theta[self.theta >= pi2] = pi2
        self.theta[self.theta <= 0] = self.epsilon

        self.calc_sin_cos_tan()
        self.calc_x()
        
if __name__ == "__main__":

    x_true = np.array([0.1, 0.2, 0.3, 0.4])
    # Optimise [x1, x2, x3, x4] under the condition x1+x2+x3+x4=1, x1>0, x2>0, x3>0, and x4>4.
    x = np.array([0.25 for _ in range(4)])
    inputs = np.random.random(size=(100, 4))
    y = np.sum(inputs * x_true, axis=1) + np.random.normal(scale=0.05, size=100)
    polar = PolarDiff(4, x=x)
    for i in range(100):
        df = np.sum((np.sum(inputs * x, axis=1) - y)[:, np.newaxis] * inputs, axis=0)
        # Update the angles by up to 10/iter degrees.
        polar.update(df, 10 / (i+1))
        x = polar.x.copy()
        print("iter: {}, x_pred: {}, x_sum: {:.3f}".format(i, x, np.sum(x)))