source: sasmodels/sasmodels/kernelpy.py @ a34b811

ticket-1257-vesicle-productticket_1156ticket_822_more_unit_tests
Last change on this file since a34b811 was a34b811, checked in by Paul Kienzle <pkienzle@…>, 3 months ago

use radius_effective/radius_effective_mode/radius_effective_modes consistently throughout the code

  • Property mode set to 100644
File size: 12.8 KB
Line 
1"""
2Python driver for python kernels
3
4Calls the kernel with a vector of $q$ values for a single parameter set.
5Polydispersity is supported by looping over different parameter sets and
6summing the results.  The interface to :class:`PyModel` matches those for
7:class:`kernelcl.GpuModel` and :class:`kerneldll.DllModel`.
8"""
9from __future__ import division, print_function
10
11import logging
12
13import numpy as np  # type: ignore
14from numpy import pi
15try:
16    from numpy import cbrt
17except ImportError:
18    def cbrt(x):
19        """Return cubed root of x."""
20        return x ** (1.0/3.0)
21
22from .generate import F64
23from .kernel import KernelModel, Kernel
24
25# pylint: disable=unused-import
26try:
27    from typing import Union, Callable
28except ImportError:
29    pass
30else:
31    from . import details
32    DType = Union[None, str, np.dtype]
33# pylint: enable=unused-import
34
35logger = logging.getLogger(__name__)
36
37
38class PyModel(KernelModel):
39    """
40    Wrapper for pure python models.
41    """
42    def __init__(self, model_info):
43        # Make sure Iq is available and vectorized.
44        _create_default_functions(model_info)
45        self.info = model_info
46        self.dtype = np.dtype('d')
47        logger.info("make python model %s", self.info.name)
48
49    def make_kernel(self, q_vectors):
50        """Instantiate the python kernel with input *q_vectors*"""
51        q_input = PyInput(q_vectors, dtype=F64)
52        return PyKernel(self.info, q_input)
53
54    def release(self):
55        """
56        Free resources associated with the model.
57        """
58        pass
59
60
61class PyInput(object):
62    """
63    Make q data available to the gpu.
64
65    *q_vectors* is a list of q vectors, which will be *[q]* for 1-D data,
66    and *[qx, qy]* for 2-D data.  Internally, the vectors will be reallocated
67    to get the best performance on OpenCL, which may involve shifting and
68    stretching the array to better match the memory architecture.  Additional
69    points will be evaluated with *q=1e-3*.
70
71    *dtype* is the data type for the q vectors. The data type should be
72    set to match that of the kernel, which is an attribute of
73    :class:`GpuProgram`.  Note that not all kernels support double
74    precision, so even if the program was created for double precision,
75    the *GpuProgram.dtype* may be single precision.
76
77    Call :meth:`release` when complete.  Even if not called directly, the
78    buffer will be released when the data object is freed.
79    """
80    def __init__(self, q_vectors, dtype):
81        self.nq = q_vectors[0].size
82        self.dtype = dtype
83        self.is_2d = (len(q_vectors) == 2)
84        if self.is_2d:
85            self.q = np.empty((self.nq, 2), dtype=dtype)
86            self.q[:, 0] = q_vectors[0]
87            self.q[:, 1] = q_vectors[1]
88        else:
89            self.q = np.empty(self.nq, dtype=dtype)
90            self.q[:self.nq] = q_vectors[0]
91
92    def release(self):
93        """
94        Free resources associated with the model inputs.
95        """
96        self.q = None
97
98
99class PyKernel(Kernel):
100    """
101    Callable SAS kernel.
102
103    *kernel* is the kernel object to call.
104
105    *model_info* is the module information
106
107    *q_input* is the DllInput q vectors at which the kernel should be
108    evaluated.
109
110    The resulting call method takes the *pars*, a list of values for
111    the fixed parameters to the kernel, and *pd_pars*, a list of (value,weight)
112    vectors for the polydisperse parameters.  *cutoff* determines the
113    integration limits: any points with combined weight less than *cutoff*
114    will not be calculated.
115
116    Call :meth:`release` when done with the kernel instance.
117    """
118    def __init__(self, model_info, q_input):
119        # type: (callable, ModelInfo, List[np.ndarray]) -> None
120        self.dtype = np.dtype('d')
121        self.info = model_info
122        self.q_input = q_input
123        self.res = np.empty(q_input.nq, q_input.dtype)
124        self.dim = '2d' if q_input.is_2d else '1d'
125
126        partable = model_info.parameters
127        #kernel_parameters = (partable.iqxy_parameters if q_input.is_2d
128        #                     else partable.iq_parameters)
129        kernel_parameters = partable.iq_parameters
130        volume_parameters = partable.form_volume_parameters
131
132        # Create an array to hold the parameter values.  There will be a
133        # single array whose values are updated as the calculator goes
134        # through the loop.  Arguments to the kernel and volume functions
135        # will use views into this vector, relying on the fact that a
136        # an array of no dimensions acts like a scalar.
137        parameter_vector = np.empty(len(partable.call_parameters)-2, 'd')
138
139        # Create views into the array to hold the arguments.
140        offset = 0
141        kernel_args, volume_args = [], []
142        for p in partable.kernel_parameters:
143            if p.length == 1:
144                # Scalar values are length 1 vectors with no dimensions.
145                v = parameter_vector[offset:offset+1].reshape(())
146            else:
147                # Vector values are simple views.
148                v = parameter_vector[offset:offset+p.length]
149            offset += p.length
150            if p in kernel_parameters:
151                kernel_args.append(v)
152            if p in volume_parameters:
153                volume_args.append(v)
154
155        # Hold on to the parameter vector so we can use it to call kernel later.
156        # This may also be required to preserve the views into the vector.
157        self._parameter_vector = parameter_vector
158
159        # Generate a closure which calls the kernel with the views into the
160        # parameter array.
161        if q_input.is_2d:
162            form = model_info.Iqxy
163            qx, qy = q_input.q[:, 0], q_input.q[:, 1]
164            self._form = lambda: form(qx, qy, *kernel_args)
165        else:
166            form = model_info.Iq
167            q = q_input.q
168            self._form = lambda: form(q, *kernel_args)
169
170        # Generate a closure which calls the form_volume if it exists.
171        self._volume_args = volume_args
172        volume = model_info.form_volume
173        shell = model_info.shell_volume
174        radius = model_info.radius_effective
175        self._volume = ((lambda: (shell(*volume_args), volume(*volume_args))) if shell and volume
176                        else (lambda: [volume(*volume_args)]*2) if volume
177                        else (lambda: (1.0, 1.0)))
178        self._radius = ((lambda mode: radius(mode, *volume_args)) if radius
179                        else (lambda mode: cbrt(0.75/pi*volume(*volume_args))) if volume
180                        else (lambda mode: 1.0))
181
182    def _call_kernel(self, call_details, values, cutoff, magnetic, radius_effective_mode):
183        # type: (CallDetails, np.ndarray, np.ndarray, float, bool) -> np.ndarray
184        if magnetic:
185            raise NotImplementedError("Magnetism not implemented for pure python models")
186        #print("Calling python kernel")
187        #call_details.show(values)
188        radius = ((lambda: 0.0) if radius_effective_mode == 0
189                  else (lambda: self._radius(radius_effective_mode)))
190        self.result = _loops(
191            self._parameter_vector, self._form, self._volume, radius,
192            self.q_input.nq, call_details, values, cutoff)
193
194    def release(self):
195        # type: () -> None
196        """
197        Free resources associated with the kernel.
198        """
199        self.q_input.release()
200        self.q_input = None
201
202
203def _loops(parameters,    # type: np.ndarray
204           form,          # type: Callable[[], np.ndarray]
205           form_volume,   # type: Callable[[], float]
206           form_radius,   # type: Callable[[], float]
207           nq,            # type: int
208           call_details,  # type: details.CallDetails
209           values,        # type: np.ndarray
210           cutoff,        # type: float
211          ):
212    # type: (...) -> None
213    ################################################################
214    #                                                              #
215    #   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!   #
216    #   !!                                                    !!   #
217    #   !!  KEEP THIS CODE CONSISTENT WITH KERNEL_TEMPLATE.C  !!   #
218    #   !!                                                    !!   #
219    #   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!   #
220    #                                                              #
221    ################################################################
222
223    # WARNING: Trickery ahead
224    # The parameters[] vector is embedded in the closures for form(),
225    # form_volume() and form_radius().  We set the initial vector from
226    # the values for the model parameters. As we loop through the polydispesity
227    # mesh, we update the components with the polydispersity values before
228    # calling the respective functions.
229    n_pars = len(parameters)
230    parameters[:] = values[2:n_pars+2]
231
232    if call_details.num_active == 0:
233        total = form()
234        weight_norm = 1.0
235        weighted_shell, weighted_form = form_volume()
236        weighted_radius = form_radius()
237
238    else:
239        pd_value = values[2+n_pars:2+n_pars + call_details.num_weights]
240        pd_weight = values[2+n_pars + call_details.num_weights:]
241
242        weight_norm = 0.0
243        weighted_form = 0.0
244        weighted_shell = 0.0
245        weighted_radius = 0.0
246        partial_weight = np.NaN
247        weight = np.NaN
248
249        p0_par = call_details.pd_par[0]
250        p0_length = call_details.pd_length[0]
251        p0_index = p0_length
252        p0_offset = call_details.pd_offset[0]
253
254        pd_par = call_details.pd_par[:call_details.num_active]
255        pd_offset = call_details.pd_offset[:call_details.num_active]
256        pd_stride = call_details.pd_stride[:call_details.num_active]
257        pd_length = call_details.pd_length[:call_details.num_active]
258
259        total = np.zeros(nq, 'd')
260        for loop_index in range(call_details.num_eval):
261            # Update polydispersity parameter values.
262            if p0_index == p0_length:
263                pd_index = (loop_index//pd_stride)%pd_length
264                parameters[pd_par] = pd_value[pd_offset+pd_index]
265                partial_weight = np.prod(pd_weight[pd_offset+pd_index][1:])
266                p0_index = loop_index%p0_length
267
268            weight = partial_weight * pd_weight[p0_offset + p0_index]
269            parameters[p0_par] = pd_value[p0_offset + p0_index]
270            p0_index += 1
271            if weight > cutoff:
272                # Call the scattering function.
273                # Assume that NaNs are only generated if the parameters are bad;
274                # exclude all q for that NaN.  Even better would be to have an
275                # INVALID expression like the C models, but that is expensive.
276                Iq = np.asarray(form(), 'd')
277                if np.isnan(Iq).any():
278                    continue
279
280                # Update value and norm.
281                total += weight * Iq
282                weight_norm += weight
283                shell, form = form_volume()
284                weighted_form += weight * form
285                weighted_shell += weight * shell
286                weighted_radius += weight * form_radius()
287
288    result = np.hstack((total, weight_norm, weighted_form, weighted_shell, weighted_radius))
289    return result
290
291
292def _create_default_functions(model_info):
293    """
294    Autogenerate missing functions, such as Iqxy from Iq.
295
296    This only works for Iqxy when Iq is written in python. :func:`make_source`
297    performs a similar role for Iq written in C.  This also vectorizes
298    any functions that are not already marked as vectorized.
299    """
300    # Note: Must call create_vector_Iq before create_vector_Iqxy.
301    _create_vector_Iq(model_info)
302    _create_vector_Iqxy(model_info)
303
304
305def _create_vector_Iq(model_info):
306    """
307    Define Iq as a vector function if it exists.
308    """
309    Iq = model_info.Iq
310    if callable(Iq) and not getattr(Iq, 'vectorized', False):
311        #print("vectorizing Iq")
312        def vector_Iq(q, *args):
313            """
314            Vectorized 1D kernel.
315            """
316            return np.array([Iq(qi, *args) for qi in q])
317        vector_Iq.vectorized = True
318        model_info.Iq = vector_Iq
319
320
321def _create_vector_Iqxy(model_info):
322    """
323    Define Iqxy as a vector function if it exists, or default it from Iq().
324    """
325    Iqxy = getattr(model_info, 'Iqxy', None)
326    if callable(Iqxy):
327        if not getattr(Iqxy, 'vectorized', False):
328            #print("vectorizing Iqxy")
329            def vector_Iqxy(qx, qy, *args):
330                """
331                Vectorized 2D kernel.
332                """
333                return np.array([Iqxy(qxi, qyi, *args) for qxi, qyi in zip(qx, qy)])
334            vector_Iqxy.vectorized = True
335            model_info.Iqxy = vector_Iqxy
336    else:
337        #print("defaulting Iqxy")
338        # Iq is vectorized because create_vector_Iq was already called.
339        Iq = model_info.Iq
340        def default_Iqxy(qx, qy, *args):
341            """
342            Default 2D kernel.
343            """
344            return Iq(np.sqrt(qx**2 + qy**2), *args)
345        default_Iqxy.vectorized = True
346        model_info.Iqxy = default_Iqxy
Note: See TracBrowser for help on using the repository browser.