<commit_before>import numpy as np
import torch
from torch.autograd import Variable


def fit(func, parameters, observations, iter=1000, lr=0.1):
    """
    Estimates the parameters of an arbitrary function via maximum likelihood estimation and
    uses plain old gradient descent for optimization
    :param func: Lambda function
    :param parameters:
    :param observations: Observations from an unknown pdf which parameters we want to estimate
    :param iter: Maximum number of iterations
    :param lr: Learning rate
    :return:
    """

    for i in range(iter):
        # Define objective function (log-likelihood) to maximize
        likelihood = torch.mean(torch.log(func(observations)))

        # Determine gradients
        likelihood.backward()

        # Update parameters with gradient descent
        for param in parameters:
            param.data.add_(lr * param.grad.data)
            param.grad.data.zero_()


def normal_dist_sample(x, mean, std):
    mean = mean.expand(x.size())
    var = std.expand(x.size()) ** 2
    return 1./torch.sqrt(2.0*np.pi*var) * torch.exp(- ((x-mean)**2) / (2.0 * var))


if __name__ == '__main__':
    """
    Estimate mean and std of a normal distribution via MLE on 10000 observations
    """
    dtype = torch.cuda.DoubleTensor

    mean = Variable(torch.zeros(1).type(dtype), requires_grad=True)
    std = Variable(torch.ones(1).type(dtype), requires_grad=True)

    func = lambda x: normal_dist_sample(x, mean, std)

    # Sample observations from a normal distribution function with different parameter
    true_mean = np.random.rand()
    true_std = np.random.rand()
    x = true_mean + np.random.randn(10000) * true_std
    x = Variable(torch.from_numpy(x)).type(dtype)

    fit(func, [mean, std], x)

    print('Estimated parameter: {{{}, {}}}, True parameter: {{{}, {}}}'.format(mean.data[0], std.data[0], true_mean, true_std))
<commit_msg>switch to numpy docstring<commit_after>import numpy as np
import torch
from torch.autograd import Variable


def fit(func, parameters, observations, iter=1000, lr=0.1):
    """Estimates the parameters of an arbitrary function via maximum likelihood estimation and
    uses plain old gradient descent for optimization


    Parameters
    ----------
    func :          Callable pdf
                    Callable probability density function (likelihood function)
                    expecting an array of observations as the only argument.
    parameters :    List
                    List of parameters that are subject to optimization.
    observations :  ndarray
                    Observations from an unknown pdf which parameters are subject to be estimated
    iter :          float
                    Maximum number of iterations
    lr :            float
                    Gradient descent learning rate

    Returns
    -------

    """

    for i in range(iter):
        # Define objective function (log-likelihood) to maximize
        likelihood = torch.mean(torch.log(func(observations)))

        # Determine gradients
        likelihood.backward()

        # Update parameters with gradient descent
        for param in parameters:
            param.data.add_(lr * param.grad.data)
            param.grad.data.zero_()


def normal_pdf(x, mean, std):
    mean = mean.expand(x.size())
    var = std.expand(x.size()) ** 2
    return 1./torch.sqrt(2.0*np.pi*var) * torch.exp(- ((x-mean)**2) / (2.0 * var))


if __name__ == '__main__':
    """
    Estimate mean and std of a normal distribution via MLE on 10000 observations
    """
    dtype = torch.cuda.DoubleTensor

    mean = Variable(torch.zeros(1).type(dtype), requires_grad=True)
    std = Variable(torch.ones(1).type(dtype), requires_grad=True)

    func = lambda x: normal_pdf(x, mean, std)

    # Sample observations from a normal distribution function with different parameter
    true_mean = np.random.rand()
    true_std = np.random.rand()
    x = true_mean + np.random.randn(10000) * true_std
    x = Variable(torch.from_numpy(x)).type(dtype)

    fit(func, [mean, std], x)

    print('Estimated parameter: {{{}, {}}}, True parameter: {{{}, {}}}'.format(mean.data[0], std.data[0], true_mean, true_std))

{'fraction_non_alphanumeric': 0.11137667304015296, 'fraction_numerical': 0.012906309751434034, 'mean_word_length': 2.797640653357532, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1489961', 'n_tokens_mistral': 1237, 'n_tokens_neox': 1176, 'n_words': 367}