Source code for xicsrt.tools.xicsrt_spread

# -*- coding: utf-8 -*-
"""
.. Authors
    Novimir Pablant <npablant@pppl.gov>
    James Kring <jdk0026@tigermail.auburn.edu>
    Yevgeniy Yakusevich <eugenethree@gmail.com>

A set of algorithms to generate vector distributions.

.. Note::
    The term 'spread' here denotes the angular range of emission. The term
    'divergence' is not used because I normally think of a divergence as a
    gaussian distributed probability distribution of angles. This type of
    distribution is available, but for generality and consistency 'spread'
    will be used throughout.
"""

import numpy as np


[docs] def vector_distribution(spread, number, name=None): """ A convenience function to retrieve vector distributions by name. Parameters ---------- spread : float or array [radians] Can be a scalar or an array. See individual distributions for format and definitions. number : int The number of vectors to generate. name : string ('isotropic') The name of the vector distribution. Available names: 'isotropic', 'isotropic_xy', 'flat', 'flat_xy', 'gaussian'. Returns ------- ndarray A numpy array of shape (number, 3) containing the generated unit vectors. """ if name is None: name = 'isotropic' name = name.lower() if name == 'isotropic': func = vector_dist_isotropic elif name == 'isotropic_xy': func = vector_dist_isotropic_xy elif name == 'flat': func = vector_dist_flat elif name == 'flat_xy': func = vector_dist_flat_xy elif name == 'gaussian': func = vector_dist_gaussian else: raise Exception(f'Distribution "{name}" is not known.') return func(spread, number)
[docs] def solid_angle(spread, name=None): """ A convenience function to retrieve solid angles that correspond to the various vector distributions. Units: [sr] """ if name is None: name = 'isotropic' name = name.lower() if name == 'isotropic': func = solid_angle_isotropic elif name == 'isotropic_xy': func = solid_angle_isotropic_xy else: raise Exception(f'Solid angle calculation for "{name}" is not available.') return func(spread)
[docs] def vector_dist_isotropic(spread, number): """ Return unit vectors from an isotropic (uniform spherical) distribution that fall within an angular spread (divergence) of theta. The ray cone is aligned along the z-axis. Parameters ---------- spread : float [radians] The half-angle of the emitted cone of vectors (axis to edge). number : int The number of vectors to generate. Returns ------- ndarray A numpy array of shape (number, 3) containing the generated unit vectors. """ theta = _parse_spread_single(spread) z = np.random.uniform(np.cos(theta), 1, number) phi = np.random.uniform(0, 2*np.pi, number) output = np.empty((number, 3)) output[:, 0] = np.sqrt(1 - z**2) * np.cos(phi) output[:, 1] = np.sqrt(1 - z**2) * np.sin(phi) output[:, 2] = z return output
[docs] def solid_angle_isotropic(spread): """ Calculate the solid angle for the vector_dist_isotropic distribution. Parameters ---------- spread : float [radians] The half-angle of cone of vectors (axis to edge). Returns ------- solid_angle Units: [sr] """ theta = _parse_spread_single(spread) solid_angle = 4 * np.pi * np.sin(theta[0]/2)**2 return solid_angle
[docs] def vector_dist_isotropic_xy(spread, number): """ Return random unit vectors from an isotroptic (uniform spherical) distribution that fall within a given x and y angular spread. The truncated-cone of vectors is aligned along the z-axis. .. Note:: This routine repeatedly filters from a circular distribution, which is accurate but not efficent. Efficiency goes down for more unequal values of the x and y spread. .. Todo:: Replace vector_dist_isotropic_xy with a more efficent calculation. A possible approach is to calculate the 2D Joint Cumulative Distribution Function for isotropic emission on a flat plane. Parameters ---------- spread : float or array [radians] | The half-angles in the x and y directions that define the extent of the | truncated-cone of vectors. Spread can be contain either 1,2 or 4 values. | | s or [s] | A single value that will be used for both the x and y directions. | [x, y] | Two values values that will be used for the x and y directions. | [xmin, xmax, ymin, ymax] | For values that define the asymmetric exent in x and y directions. | Example: [-0.1, 0.1, -0.5, 0.5] number : int The number of vectors to generate. Returns ------- ndarray A numpy array of shape (number, 3) containing the generated unit vectors. """ theta = _parse_spread_xy(spread) # Determine the extent of the circular cone than we need. theta_xmax = np.max(np.abs(theta[0:2])) theta_ymax = np.max(np.abs(theta[2:])) theta_max = np.arcsin(np.sqrt(np.sin(theta_xmax) ** 2 + np.sin(theta_ymax) ** 2)) # Generate and filter rays until we have the requested number. output = np.empty((number, 3)) n_filled = 0 while n_filled < number: vectors = vector_dist_isotropic(theta_max, number) mask = np.ones(number, dtype=bool) mask &= vectors[:, 0] / np.sqrt(vectors[:, 0] ** 2 + vectors[:, 2] ** 2) > np.sin(theta[0]) mask &= vectors[:, 0] / np.sqrt(vectors[:, 0] ** 2 + vectors[:, 2] ** 2) <= np.sin(theta[1]) mask &= vectors[:, 1] / np.sqrt(vectors[:, 1] ** 2 + vectors[:, 2] ** 2) > np.sin(theta[2]) mask &= vectors[:, 1] / np.sqrt(vectors[:, 1] ** 2 + vectors[:, 2] ** 2) <= np.sin(theta[3]) n_new = np.sum(mask) n_need = number - n_filled if n_new > (number - n_filled): n_new = n_need output[n_filled:n_filled + n_new, :] = vectors[mask, :][:n_new, :] n_filled += n_new return output
[docs] def solid_angle_isotropic_xy(spread): """ Calculate the solid angle for the vector_dist_isotropic_xy distribution. Units: [sr] """ theta = _parse_spread_xy(spread) solid_angle = ( np.arcsin(np.abs(np.sin(theta[0])*np.sin(theta[2]))) + np.arcsin(np.abs(np.sin(theta[0])*np.sin(theta[3]))) + np.arcsin(np.abs(np.sin(theta[1])*np.sin(theta[2]))) + np.arcsin(np.abs(np.sin(theta[1])*np.sin(theta[3]))) ) return solid_angle
[docs] def vector_dist_flat(spread, number): """ Return unit vectors from an flat (uniform planar) distribution that fall within an angular spread. The ray cone is aligned along the z-axis. Parameters ---------- spread : float [radians] The half-angle of the emitted cone of vectors (axis to edge). number : int The number of vectors to generate. Returns ------- ndarray A numpy array of shape (number, 3) containing the generated unit vectors. """ theta = _parse_spread_single(spread) r = np.sqrt(np.random.uniform(0, np.tan(theta), number)) angle1 = np.random.uniform(0, 2*np.pi, number) angle0 = np.arctan(r) output = np.empty((number, 3)) output[:, 0] = np.cos(angle1) * np.sin(angle0) output[:, 1] = np.sin(angle1) * np.sin(angle0) output[:, 2] = np.cos(angle0) return output
[docs] def vector_dist_flat_xy(spread, number): """ Return random unit vectors from an flat (uniform planar) distribution that fall within a given x and y angular spread. The truncated-cone of vectors is aligned along the z-axis. .. Note:: This distribution is identical to that used by the SHADOW raytracing code for both the 'flat' and 'uniform' distributions (as of 2021-01). Parameters ---------- spread : float or array [radians] | The half-angles in the x and y directions that define the extent of the | truncated-cone of vectors. Spread can be contain either 1,2 or 4 values. | | s or [s] | A single value that will be used for both the x and y directions. | [x, y] | Two values values that will be used for the x and y directions. | [xmin, xmax, ymin, ymax] | For values that define the asymmetric exent in x and y directions. | Example: [-0.1, 0.1, -0.5, 0.5] number : int The number of vectors to generate. Returns ------- ndarray A numpy array of shape (number, 3) containing the generated unit vectors. """ theta = _parse_spread_xy(spread) range = np.tan(theta) x = np.random.uniform(range[0], range[1], number) y = np.random.uniform(range[2], range[3], number) angle0 = np.arctan(np.sqrt(x ** 2 + y ** 2)) angle1 = np.arctan2(y, x) output = np.empty((number, 3)) output[:, 0] = np.cos(angle1) * np.sin(angle0) output[:, 1] = np.sin(angle1) * np.sin(angle0) output[:, 2] = np.cos(angle0) return output
[docs] def vector_dist_flat_gaussian(spread, num_samples): """ Create distribution of vectors with a Gaussian distribution on a flat plane. The ray code is aligned with the z-axis, so the distribution is Gaussian in the x-y directions. For small angles this distribution will approximate the Kent distribution and will be approximately gaussian in angle. Parameters ---------- spread : float [radians] The half-with-at-half-max (hwhm) of the Gaussian angular distribution. Returns ------- ndarray A numpy array of shape (number, 3) containing the generated unit vectors. """ theta = _parse_spread_single(spread) # Convert the angluar hwhm to sigma. sigma = theta[0] / (np.sqrt(2 * np.log(2))) # Calculate the equivilant width in x & y on a plane with z=1. xsigma = np.sin(sigma) ysigma = np.sin(sigma) # Origins for a gaussian distribution of rays. mean = [0, 0] # Define a diagnonal covariance matrix. cov = [[xsigma**2, 0], [0, ysigma**2]] x, y = np.random.multivariate_normal(mean, cov, num_samples).T z = np.full(num_samples, 1.0) out = np.vstack((x, y, z)).T out_mod = np.linalg.norm(out, axis=1) out = np.einsum('ij,i->ij', out, 1 / np.linalg.norm(out, axis=1), optimize=True) return out
[docs] def _to_ndarray(spread): """ Convert input value to an numpy array. Scalars will be transformed to a one element array. """ if np.isscalar(spread): out = np.array([spread]) elif isinstance(spread, np.ndarray): out = spread else: out = np.array(spread) return out
[docs] def _parse_spread_single(spread): """ Parse the number of input values in spread and return a standard array. Use only for distribution with a single spread value. """ spread = _to_ndarray(spread) if len(spread) != 1: raise Exception('Spread must be a scalar or one element array.') return spread
[docs] def _parse_spread_xy(spread): """ Parse the number of input values in spread and return a standard array. Use only for assymetric xy distributions. """ spread = _to_ndarray(spread) if len(spread) == 1: out = [-spread[0], spread[0], -spread[0], spread[0]] elif len(spread) == 2: out = [-spread[0], spread[0], -spread[1], spread[1]] elif len(spread) == 4: out = [spread[0], spread[1], spread[2], spread[3]] else: raise Exception('Spread must have 1, 2 or 3 elements. See docstring.') return out