Module ukat.mapping.fitting.relaxation

class Model:
    def __init__(self, pixel_array, x, eq, mask=None, multithread=True):
        """
        A template class for fitting models to pixel arrays.

        Parameters
        ----------
        pixel_array : np.ndarray
            An array containing the signal from each voxel with the last
            dimension being the dependent variable axis
        x : np.ndarray
            An array containing the dependent variable e.g. time
        eq : function
            A function that takes the dependent variable as the first
            argument and the parameters to fit as the remaining arguments
        mask : np.ndarray, optional
            A boolean mask of voxels to fit. Should be the shape of the desired
            map rather than the raw data i.e. omit the dependent variable axis
        multithread : bool, optional
            Default True
            If True, the fitting will be performed in parallel using all
            available cores
        """
        # Attributes that can be set from default inputs
        self.pixel_array = pixel_array
        self.map_shape = pixel_array.shape[:-1]
        self.x = x
        self.eq = eq
        self.mask = mask
        self.multithread = multithread
        self.n_x = pixel_array.shape[-1]
        self.n_params = self._get_n_params()

        # Placeholder attributes that will be overwritten by the child class
        self.initial_guess = None
        self.signal_list = None
        self.x_list = None
        self.p0_list = None
        self.mask_list = None

    def generate_lists(self):
        """
        Generate the lists of data, dependent variables, initial guesses and
        masks to be used in the fitting process
        """
        self.signal_list = self.pixel_array.reshape(-1, self.n_x).tolist()
        self.x_list = [self.x] * len(self.signal_list)
        self.p0_list = [self.initial_guess] * len(self.signal_list)
        self.mask_list = self._get_mask_list()

    def _get_n_params(self):
        """
        Get the number of parameters to fit

        Returns
        -------
        n_params : int
            The number of parameters to fit
        """
        n_params = len(inspect.signature(self.eq).parameters) - 1
        return n_params

    def _get_mask_list(self):
        """
        Get a list of masks to be used in the fitting process, if no mask
        has been specified it will be a list of True i.e. all voxels will be
        fit

        Returns
        -------
        mask_list : list
            A list of booleans indicating whether to fit a voxel or not
        """
        if self.mask is None:
            mask_list = [True] * len(self.signal_list)
            return mask_list
        else:
            mask_list = self.mask.reshape(-1).tolist()
            return mask_list