caustics package

Subpackages

Submodules

caustics.constants module

caustics.cosmology module

class caustics.cosmology.Cosmology(name: str | None = None)[source]

Bases: Parametrized

Abstract base class for cosmological models.

This class provides an interface for cosmological computations used in lensing such as comoving distance and critical surface density.

Units:

  • Distance: Mpc

  • Mass: solar mass

name

Name of the cosmological model.

Type:

str

angular_diameter_distance(z: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the angular diameter distance to redshift z.

Parameters:

  • z (Tensor) – The redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The angular diameter distance to each redshift.

Return type:

Tensor

angular_diameter_distance_z1z2(z1: Tensor, z2: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the angular diameter distance between two redshifts.

Parameters:

  • z1 (Tensor) – The starting redshifts.

  • z2 (Tensor) – The ending redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The angular diameter distance between each pair of redshifts.

Return type:

Tensor

abstract comoving_distance(z: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the comoving distance to redshift z.

Parameters:

  • z (Tensor) – The redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The comoving distance to each redshift.

Return type:

Tensor

comoving_distance_z1z2(z1: Tensor, z2: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the comoving distance between two redshifts.

Parameters:

  • z1 (Tensor) – The starting redshifts.

  • z2 (Tensor) – The ending redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The comoving distance between each pair of redshifts.

Return type:

Tensor

abstract critical_density(z: Tensor, params: Packed | None = None) → Tensor[source]

Compute the critical density at redshift z.

Parameters:

  • z (Tensor) – The redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The critical density at each redshift.

Return type:

Tensor

critical_surface_density(z_l: Tensor, z_s: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the critical surface density between lens and source planes.

Parameters:

  • z_l (Tensor) – The lens redshifts.

  • z_s (Tensor) – The source redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The critical surface density for each pair of lens and source redshifts.

Return type:

Tensor

time_delay_distance(z_l: Tensor, z_s: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the time delay distance between lens and source planes.

Parameters:

  • z_l (Tensor) – The lens redshifts.

  • z_s (Tensor) – The source redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The time delay distance for each pair of lens and source redshifts.

Return type:

Tensor

abstract transverse_comoving_distance(z: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the transverse comoving distance to redshift z (Mpc).

Parameters:

  • z (Tensor) – The redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The transverse comoving distance to each redshift in Mpc.

Return type:

Tensor

transverse_comoving_distance_z1z2(z1: Tensor, z2: Tensor, *args, params: Packed | None = None) → Tensor[source]

Compute the transverse comoving distance between two redshifts (Mpc).

Parameters:

  • z1 (Tensor) – The starting redshifts.

  • z2 (Tensor) – The ending redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The transverse comoving distance between each pair of redshifts in Mpc.

Return type:

Tensor

class caustics.cosmology.FlatLambdaCDM(h0: Tensor | None = tensor(0.6766), critical_density_0: Tensor | None = tensor(1.2705e+11), Om0: Tensor | None = tensor(0.3097), name: str | None = None)[source]

Bases: Cosmology

Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) cosmology with no radiation.

comoving_distance(z: Tensor, h0, central_critical_density, Om0, *args, params: Packed | None = None) → Tensor[source]

Calculate the comoving distance to redshift z.

Parameters:

  • z (Tensor) – Redshift.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

Comoving distance to redshift z.

Return type:

Tensor

critical_density(z: Tensor, h0, central_critical_density, Om0, *args, params: Packed | None = None) → Tensor[source]

Calculate the critical density at redshift z.

Parameters:

  • z (Tensor) – Redshift.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

Critical density at redshift z.

Return type:

torch.Tensor

hubble_distance(h0)[source]

Calculate the Hubble distance.

Parameters:

h0 (Tensor) – Hubble constant.

Returns:

Hubble distance.

Return type:

Tensor

to(device: device | None = None, dtype: dtype | None = None)[source]

Moves static Params for this component and its childs to the specified device and casts them to the specified data type.

transverse_comoving_distance(z: Tensor, h0, central_critical_density, Om0, *args, params: Packed | None = None) → Tensor[source]

Compute the transverse comoving distance to redshift z (Mpc).

Parameters:

  • z (Tensor) – The redshifts.

  • params (Packed, optional) – Dynamic parameter container for the computation.

Returns:

The transverse comoving distance to each redshift in Mpc.

Return type:

Tensor

caustics.namespace_dict module

class caustics.namespace_dict.NamespaceDict[source]

Bases: OrderedDict

Add support for attributes on top of an OrderedDict

class caustics.namespace_dict.NestedNamespaceDict[source]

Bases: _NestedNamespaceDict

Example usage ```python

nested_namespace = NestedNamespaceDict() nested_namespace.foo = ‘Hello’ nested_namespace.bar = {‘baz’: ‘World’} nested_namespace.bar.qux = 42 # works also in the follwoing way

nested_namespace[“bar.qux”] = 42

print(nested_namespace) # Output: # {‘foo’: ‘Hello’, ‘bar’: {‘baz’: ‘World’, ‘qux’: 42 }}

#============================== # Flattened key access #============================== print(nested_dict[‘foo’]) # Output: Hello print(nested_dict[‘bar.baz’]) # Output: World print(nested_dict[‘bar.qux’]) # Output: 42

#============================== # Nested namespace access #============================== print(nested_dict.bar.qux) # Output: 42

#============================== # Flatten and collapse method #============================== print(nested_dict.flatten()) # Output: # {‘foo’: ‘Hello’, ‘bar.baz’: ‘World’, ‘bar.qux’: 42}

print(nested_dict.collapse() # Output: # {‘foo’: ‘Hello’, ‘baz’: ‘World’, ‘qux’: 42}

caustics.packed module

class caustics.packed.Packed[source]

Bases: OrderedDict

Dummy wrapper for x so other functions can check its type.

caustics.parameter module

class caustics.parameter.Parameter(value: Tensor | float | None = None, shape: tuple[int, ...] | None = ())[source]

Bases: object

Represents a static or dynamic parameter used for strong gravitational lensing simulations in the caustics codebase.

A static parameter has a fixed value, while a dynamic parameter must be passed in each time it’s required.

value

The value of the parameter.

Type:

Optional[Tensor]

shape

The shape of the parameter.

Type:

tuple[int, …]

property dtype
property dynamic: bool
set_static()[source]
property shape: tuple[int, ...]
property static: bool
to(device: device | None = None, dtype: dtype | None = None)[source]

Moves and/or casts the values of the parameter.

Parameters:

  • device (Optional[torch.device], optional) – The device to move the values to. Defaults to None.

  • dtype (Optional[torch.dtype], optional) – The desired data type. Defaults to None.

property value: Tensor | None

caustics.parametrized module

class caustics.parametrized.Parametrized(name: str | None = None)[source]

Bases: object

Represents a class with Param and Parametrized attributes, typically used to construct parts of a simulator that have parameters which need to be tracked during MCMC sampling.

This class can contain Params, Parametrized, tensor buffers or normal attributes as its attributes. It provides functionalities to manage these attributes, ensuring that an attribute of one type isn’t rebound to be of a different type.

TODO - Attributes can be Params, Parametrized, tensor buffers or just normal attributes. - Need to make sure an attribute of one of those types isn’t rebound to be of a different type.

name

The name of the Parametrized object. Default to class name.

Type:

str

parents

Nested dictionary of parent Parametrized objects (higher level, more abstract modules).

Type:

NestedNamespaceDict

params

Dictionary of parameters.

Type:

OrderedDict[str, Parameter]

childs NestedNamespaceDict

Nested dictionary of childs Parametrized objects (lower level, more specialized modules).

dynamic_size

Size of dynamic parameters.

Type:

int

n_dynamic

Number of dynamic parameters.

Type:

int

n_static

Number of static parameters.

Type:

int

add_param(name: str, value: Tensor | float | None = None, shape: tuple[int, ...] | None = ())[source]

Stores a parameter in the _params dictionary and records its size.

Parameters:

  • name (str) – The name of the parameter.

  • value (Optional[Tensor], optional) – The value of the parameter. Defaults to None.

  • shape (Optional[tuple[int, ...]], optional) – The shape of the parameter. Defaults to an empty tuple.

add_parametrized(p: Parametrized, set_attr=True)[source]

Add a child to this module, and create edges for the DAG

property dynamic_modules: NamespaceDict[str, Parametrized]
property dynamic_size: int
get_graph(show_dynamic_params: bool = False, show_static_params: bool = False) → graphviz.Digraph[source]

Returns a graph representation of the object and its parameters.

Parameters:

  • show_dynamic_params (bool, optional) – If true, the dynamic parameters are shown in the graph. Defaults to False.

  • show_static_params (bool, optional) – If true, the static parameters are shown in the graph. Defaults to False.

Returns:

The graph representation of the object.

Return type:

graphviz.Digraph

property module_params: NestedNamespaceDict
property n_dynamic: int
property n_static: int
property name: str
pack(x: list[Tensor] | dict[str, list[Tensor] | Tensor | dict[str, Tensor]] | Tensor = {}) → Packed[source]

Converts a list or tensor into a dict that can subsequently be unpacked into arguments to this component and its childs. Also, add a batch dimension to each Tensor without such a dimension.

Parameters:

x (Union[list[Tensor], dict[str, Union[list[Tensor], Tensor, dict[str, Tensor]]], Tensor) – The input to be packed. Can be a list of tensors, a dictionary of tensors, or a single tensor.

Returns:

The packed input, and whether or not the input was batched.

Return type:

Packed

Raises:

  • ValueError – If the input is not a list, dictionary, or tensor.

  • ValueError – If the input is a dictionary and some keys are missing.

  • ValueError – If the number of dynamic arguments does not match the expected number.

  • ValueError – If the input is a tensor and the shape does not match the expected shape.

property params: NestedNamespaceDict
to(device: device | None = None, dtype: dtype | None = None)[source]

Moves static Params for this component and its childs to the specified device and casts them to the specified data type.

unpack(x: dict[str, list[Tensor] | dict[str, Tensor] | Tensor] | None) → list[Tensor][source]

Unpacks a dict of kwargs, list of args or flattened vector of args to retrieve this object’s static and dynamic parameters.

Parameters:

x (Optional[dict[str, Union[list[Tensor], dict[str, Tensor], Tensor]]]) – The packed object to be unpacked.

Returns:

Unpacked static and dynamic parameters of the object. Note that parameters will have an added batch dimension from the pack method.

Return type:

list[Tensor]

Raises:

  • ValueError – If the input is not a dict, list, tuple or tensor.

  • ValueError – If the argument type is invalid. It must be a dict containing key {self.name} and value containing args as list or flattened tensor, or kwargs.

caustics.parametrized.unpack(n_leading_args=0)[source]

caustics.utils module

caustics.utils.batch_lm(X, Y, f, C=None, epsilon=0.1, L=1.0, L_dn=11.0, L_up=9.0, max_iter=50, L_min=1e-09, L_max=1000000000.0, stopping=0.0001, f_args=(), f_kwargs={})[source]
caustics.utils.derotate(vx, vy, phi: Tensor | None = None)[source]

Applies inverse rotation to the velocity components (vx, vy) using the rotation angle phi.

Parameters:

  • vx (Tensor) – Tensor containing the x-component of velocity.

  • vy (Tensor) – Tensor containing the y-component of velocity.

  • phi (Optional[Tensor], optional) – Tensor containing the rotation angles. If None, no rotation is applied. Defaults to None.

Returns:

Tuple containing the derotated x and y components of velocity.

Return type:

Tuple[Tensor, Tensor]

caustics.utils.flip_axis_ratio(q, phi)[source]

Makes the value of ‘q’ positive, then swaps x and y axes if ‘q’ is larger than 1.

Parameters:

  • q (Tensor) – Tensor containing values to be processed.

  • phi (Tensor) – Tensor containing the phi values for the orientation of the axes.

Returns:

Tuple containing the processed ‘q’ and ‘phi’ Tensors.

Return type:

Tuple[Tensor, Tensor]

caustics.utils.gaussian(pixelscale, nx, ny, sigma, upsample=1, dtype=torch.float32, device=None)[source]
caustics.utils.get_cluster_means(xs: Tensor, k: int)[source]

Computes cluster means using the k-means++ initialization algorithm.

Parameters:

  • xs (Tensor) – A tensor of data points.

  • k (int) – The number of clusters.

Returns:

A tensor of cluster means.

Return type:

Tensor

caustics.utils.get_meshgrid(pixelscale, nx, ny, device=None, dtype=torch.float32) → Tuple[Tensor, Tensor][source]

Generates a 2D meshgrid based on the provided pixelscale and dimensions.

Parameters:

  • pixelscale (float) – The scale of the meshgrid in each dimension.

  • nx (int) – The number of grid points along the x-axis.

  • ny (int) – The number of grid points along the y-axis.

  • device (torch.device, optional) – The device on which to create the tensor. Defaults to None.

  • dtype (torch.dtype, optional) – The desired data type of the tensor. Defaults to torch.float32.

Returns:

The generated meshgrid as a tuple of Tensors.

Return type:

Tuple[Tensor, Tensor]

caustics.utils.interp1d(x: Tensor, y: Tensor, xs: Tensor, extend: str = 'extrapolate') → Tensor[source]

Compute the 1D cubic spline interpolation for the given data points using PyTorch.

Parameters:

  • x (Tensor) – A 1D tensor representing the x-coordinates of the known data points.

  • y (Tensor) – A 1D tensor representing the y-coordinates of the known data points.

  • xs (Tensor) – A 1D tensor representing the x-coordinates of the positions where the cubic spline function should be evaluated.

  • extend (str, optional) – The method for handling extrapolation, either “const”, “extrapolate”, or “linear”. Default is “extrapolate”. “const”: Use the value of the last known data point for extrapolation. “linear”: Use linear extrapolation based on the last two known data points. “extrapolate”: Use cubic extrapolation of data.

Returns:

A 1D tensor representing the interpolated values at the specified positions (xs).

Return type:

Tensor

caustics.utils.interp2d(im: Tensor, x: Tensor, y: Tensor, method: str = 'linear', padding_mode: str = 'zeros') → Tensor[source]

Interpolates a 2D image at specified coordinates. Similar to torch.nn.functional.grid_sample with align_corners=False.

Parameters:

  • im (Tensor) – A 2D tensor representing the image.

  • x (Tensor) – A 0D or 1D tensor of x coordinates at which to interpolate.

  • y (Tensor) – A 0D or 1D tensor of y coordinates at which to interpolate.

  • method (str, optional) – Interpolation method. Either ‘nearest’ or ‘linear’. Defaults to ‘linear’.

  • padding_mode (str, optional) – Defines the padding mode when out-of-bound indices are encountered. Either ‘zeros’ or ‘extrapolate’. Defaults to ‘zeros’.

Raises:

  • ValueError – If im is not a 2D tensor.

  • ValueError – If x is not a 0D or 1D tensor.

  • ValueError – If y is not a 0D or 1D tensor.

  • ValueError – If padding_mode is not ‘extrapolate’ or ‘zeros’.

  • ValueError – If method is not ‘nearest’ or ‘linear’.

Returns:

Tensor with the same shape as x and y containing the interpolated values.

Return type:

Tensor

caustics.utils.safe_divide(num, denom, places=7)[source]

Safely divides two tensors, returning zero where the denominator is zero.

Parameters:

  • num (Tensor) – The numerator tensor.

  • denom (Tensor) – The denominator tensor.

Returns:

The result of the division, with zero where the denominator was zero.

Return type:

Tensor

caustics.utils.safe_log(x)[source]

Safely applies the logarithm to a tensor, returning zero where the tensor is zero.

Parameters:

x (Tensor) – The input tensor.

Returns:

The result of applying the logarithm, with zero where the input was zero.

Return type:

Tensor

caustics.utils.to_elliptical(x, y, q: Tensor)[source]

Converts Cartesian coordinates to elliptical coordinates.

Parameters:

  • x (Tensor) – Tensor containing the x-coordinates.

  • y (Tensor) – Tensor containing the y-coordinates.

  • q (Tensor) – Tensor containing the elliptical parameters.

Returns:

Tuple containing the x and y coordinates in elliptical form.

Return type:

Tuple[Tensor, Tensor]

caustics.utils.translate_rotate(x, y, x0, y0, phi: Tensor | None = None)[source]

Translates and rotates the points (x, y) by subtracting (x0, y0) and applying rotation angle phi.

Parameters:

  • x (Tensor) – Tensor containing the x-coordinates.

  • y (Tensor) – Tensor containing the y-coordinates.

  • x0 (Tensor) – Tensor containing the x-coordinate translation values.

  • y0 (Tensor) – Tensor containing the y-coordinate translation values.

  • phi (Optional[Tensor], optional) – Tensor containing the rotation angles. If None, no rotation is applied. Defaults to None.

Returns:

Tuple containing the translated and rotated x and y coordinates.

Return type:

Tuple[Tensor, Tensor]

caustics.utils.vmap_n(func: Callable, depth: int = 1, in_dims: int | Tuple = 0, out_dims: int | Tuple[int, ...] = 0, randomness: str = 'error') → Callable[source]

Transforms a function depth times using torch.vmap with the same arguments passed each time. Returns func transformed depth times by vmap, with the same arguments passed to vmap each time.

Parameters:

  • func (Callable) – The function to transform.

  • depth (int, optional) – The number of times to apply torch.vmap. Defaults to 1.

  • in_dims (Union[int, Tuple], optional) – The dimensions to vectorize over in the input. Defaults to 0.

  • out_dims (Union[int, Tuple[int, ...]], optional) – The dimensions to vectorize over in the output. Defaults to 0.

  • randomness (str, optional) – How to handle randomness. Defaults to ‘error’.

Raises:

ValueError – If depth is less than 1.

Returns:

The transformed function.

Return type:

Callable

TODO: test.

Module contents