Fit sigmoid with parallel seeds to find global minimum.

fit_sigmoid.py

"""Fit sigmoid with parallel seeds to find global minimum.
"""
import matplotlib.pyplot as plt
import numpy as np
from sklearn.base import BaseEstimator
import torch
import torch.nn as nn


class SigmoidRegression(BaseEstimator):
    def __init__(self, n_epochs=5000, n_seeds=100, lr=1e-3):
        self.n_epochs = n_epochs
        self.n_seeds = n_seeds
        self.lr = lr

    def _sigmoid(self, X, a, b, c, d):
        exp = np.exp
        if isinstance(X, torch.Tensor):
            exp = torch.exp
        return a / (1 + exp(-c * (X - d))) + b

    def _checkX(self, X):
        x = np.squeeze(X)
        assert x.ndim == 1
        return x

    def _to_tensor(self, z):
        if isinstance(z, list):
            z = torch.tensor(z, dtype=torch.float)
        elif isinstance(z, np.ndarray):
            z = torch.from_numpy(z).float()
        return z

    def fit(self, X, y):
        self.loss_ = list()
        x = self._to_tensor(self._checkX(X))
        y = self._to_tensor(y)

        X = x.unsqueeze(1)
        Y = y.unsqueeze(1) * torch.ones(len(y), self.n_seeds)

        abcd = nn.Parameter(torch.randn(4, self.n_seeds)).float()
        a, b, c, d = abcd
        optimizer = torch.optim.Adam([abcd, ], lr=self.lr)
        loss = nn.MSELoss()

        for epoch in range(self.n_epochs):
            optimizer.zero_grad()
            Y_hat = self._sigmoid(X, a, b, c, d)
            loss(Y_hat, Y).backward()
            optimizer.step()

        losses = list()
        for a, b, c, d in abcd.transpose(0, 1):
            y_hat = self._sigmoid(x, a, b, c, d)
            losses.append(loss(y_hat, y).detach().numpy())
        abcd = abcd[:, np.argmin(losses)]

        self.coef_ = abcd.detach().numpy()
        return self

    def predict(self, X):
        x = self._checkX(X)
        return self._sigmoid(x, *self.coef_)


if __name__ == '__main__':

    x = [0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.]
    y = [-0.311, 0.0815, -0.309, -0.117, -0.389,
         0.525, 0.495, 0.334, 0.486, 0.791, 0.349]

    model = SigmoidRegression()
    model.fit(x, y)
    plt.scatter(x, y)
    x = np.linspace(0, 1)
    plt.plot(x, model.predict(x))