Source code for funclp.modules.make_calculation_LP.make_calculation

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Date          : 2025-12-18
# Author        : Lancelot PINCET
# GitHub        : https://github.com/LancelotPincet
# Library       : funcLP
# Module        : make_calculation

"""
Does the calculation for a given ufunc.
"""



# %% Libraries
from funclp import use_inputs, use_shapes, use_cuda, use_broadcasting



# %% Function
[docs] def make_calculation(function, name, args, kwargs, out=None, ignore=False) : ''' Does the calculation for a given ufunc. Parameters ---------- function : Function the function object where to calculate. name : str Name of the ufunc to calculate. args : tuple Argument to use for calculation. kwargs : dict Keyword arguments to use for calculation. out : xp.ndarray Output array, if None will initialize it automatically. ignore : xp.ndarray Boolean array defining for which model we ignore calculations. Returns ------- out : xp.ndarray output array. others : dict dictionnary containing all the other parameters calculated for this calculation. [nomodel, nopoint, nmodels, npoints, variables_shape, data_shape, parameters_shape, cuda, xp, transfer_back, blocks_per_grid, threads_per_block, variables, data, parameters, dtype, jitted] ''' # Get ufunc ufunc = getattr(function.__class__, name) if ufunc is None : raise SyntaxError(f'ufunc {name} was not recognized') # Adds main parameters to kwargs if hasattr(function, "parameters") : function_parameters = {key: value for key, value in function.parameters.items() if key in ufunc.parameters} kwargs = {**function_parameters, **kwargs} # Use inputs inputs = use_inputs(ufunc, args, kwargs) # variables, data, parameters # Get shapes empty_dimensions, out_shape, in_shapes = use_shapes(*inputs) # (nomodel, nopoint), (nmodels, npoints), (variables_shapes, data_shapes, parameters_shapes) variables_shape, data_shape, parameters_shape = in_shapes nmodels, npoints = out_shape nomodel, nopoint = empty_dimensions # Define cuda cuda, xp, transfer_back, blocks_per_grid, threads_per_block = use_cuda(function, out_shape, inputs) # Broadcasting variables, data, parameters, dtype = use_broadcasting(xp, *inputs, *in_shapes, out_shape) ignore = xp.broadcast_to(xp.asarray(ignore).astype(xp.bool_), parameters_shape).reshape(nmodels) # Output out = xp.empty(shape=out_shape, dtype=dtype) if out is None else xp.asarray(out) if out.shape != out_shape : raise ValueError(f'Out shape {out.shape} is does not correspond to expected {out_shape}') # Arrays constants = [xp.asarray(getattr(function, arr)) for arr in ufunc.constants] # Function jitted jitted = getattr(function, f'gpu_{name}')[blocks_per_grid, threads_per_block] if cuda else getattr(function, f'cpu_{name}') # Apply function jitted(*variables, *data, *parameters, *constants, out, ignore) # Modify output out = out.reshape(data_shape) if transfer_back : out = xp.asnumpy(out) if nomodel : out = out[0, ...] if nopoint : out = out[..., 0] # Returns others others = dict( nomodel = nomodel, nopoint = nopoint, nmodels = nmodels, npoints = npoints, variables_shape = variables_shape, data_shape = data_shape, parameters_shape = parameters_shape, cuda = cuda, xp = xp, transfer_back = transfer_back, blocks_per_grid = blocks_per_grid, threads_per_block = threads_per_block, variables = variables, data = data, parameters = parameters, dtype = dtype, jitted = jitted, ) return out, others
# %% Test function run if __name__ == "__main__": from corelp import test test(__file__)