iactrace.utils

Utility functions for sampling and optimization.

Sampling Functions

Functions for sampling points on geometric shapes:

iactrace.utils.sample_disk(key, shape)

Generate uniform random samples within a unit disk.

Args:

key: JAX random key shape: Shape of samples to generate

Returns:

2D points in unit disk (…, 2)

iactrace.utils.sample_polygon(key, vertices, shape)

Sample uniformly over a convex polygon via fan triangulation.

Args:

key: JAX random key vertices: (n, 2) array of polygon vertices in order (must be convex) shape: batch shape

Returns:

points: (…, 2) array of (x, y) coordinates

Filtering (for optimization)

Utilities for filtering pytree parameters during gradient-based optimization.

iactrace.utils.filtering.make_filter(model, trainable=None, frozen=None)

Create a filter spec for eqx.partition.

Args:

model: The pytree to filter trainable: Paths to train (everything else frozen), can be:

• str: single pattern

• Sequence[str]: list of patterns

• Callable[[path_tuple, leaf], bool]: custom filter

frozen: Paths to freeze (everything else trainable), same types

Returns:

Filter spec (pytree of bools matching model structure)

Patterns:

‘mirror_groups.rotations’ - exact match ‘mirror_groups.*’ - all direct children ‘mirror_groups.*.rotations’ - rotations in any child ‘**.rotations’ - any path ending in rotations

Examples:

# Only train mirror rotations make_filter(tel, trainable=’mirror_groups.*.rotations’)

# Train all mirror group params make_filter(tel, trainable=’mirror_groups.**’)

# Freeze only sensor positions make_filter(tel, frozen=’sensors.*.position’)

# Custom filter make_filter(tel, trainable=lambda path, leaf: ‘rotation’ in str(path))

iactrace.utils.filtering.partition(model, trainable=None, frozen=None)

Partition model into (trainable, frozen).

Args:

model: Equinox module trainable: What to train (else frozen) frozen: What to freeze (else trainable)

Returns:

(trainable_tree, frozen_tree) for use with eqx.combine

iactrace.utils.filtering.trainable_params(model, spec)

Get only trainable parameters.

iactrace.utils.filtering.show_structure(model)

Print model structure with paths - useful for debugging filter patterns.