argosim.optim_utils module

argosim.optim_utils module#

Optim utils.

This module contains the building blocks for end-to-end gradient-based optimisation of antenna-array layouts using the differentiable Kaiser-Bessel gridding forward model from argosim.imaging_utils.

It is intentionally modular so that the user can mix and match:

  • a forward model parameterised by an ObservationConfig;

  • a target / objective loss (currently: image-matching MSE; future: FWHM, sidelobe level, UV-coverage metrics);

  • an arbitrary list of weighted soft constraints (Constraint) — minimum antenna spacing, maximum array radius, centroid pin, etc.;

  • an optimiser built via make_optimizer() (Adam, SGD, AdamW, …) with an optional learning-rate schedule via exp_decay_schedule().

The full optimisation loop is exposed as optimise_array(), which returns the final positions, the loss history, and per-step antenna snapshots for animation.

Authors:

Ezequiel Centofanti <ezequiel.centofanti@cea.fr>

class argosim.optim_utils.Constraint(name: str, fn: Callable, weight: float)[source]#

Bases: object

Constraint.

A weighted soft constraint applied to the antenna array.

Parameters:

  • name (str) – Human-readable name (used in reports / logs).

  • fn (Callable) – (antenna,) -> scalar penalty function.

  • weight (float) – Multiplicative weight in the total loss.

fn: Callable#
name: str#
weight: float#
class argosim.optim_utils.ObservationConfig(npx: int = 256, fov: float = 0.4, lat: float = 0.6108652381980153, dec: float = 0.5585053606381855, freq: float = 2000000000.0, bw: float = 1000000000.0, n_freqs: int = 5, n_times: int = 5, track_h: float = 0.5, t_0: float = 1.0, kernel_W: int = 7, oversampling: int = 2)[source]#

Bases: object

Observation configuration.

Static parameters of the imaging / observation model used by the differentiable forward map. Kept frozen so it is hashable and safe to pass through jax.jit() as a static argument.

Parameters:

  • npx (int) – Number of pixels per side of the (square) image grid.

  • fov (float) – Field of view in degrees (square FOV).

  • lat (float) – Array latitude in radians.

  • dec (float) – Source declination in radians.

  • freq (float) – Central observation frequency in Hz.

  • bw (float) – Total bandwidth in Hz.

  • n_freqs (int) – Number of frequency channels sampled across bw.

  • n_times (int) – Number of time samples across the integration.

  • track_h (float) – Total integration time in hours.

  • t_0 (float) – Initial hour-angle in hours.

  • kernel_W (int) – KB gridding kernel support width in pixels (odd integer).

  • oversampling (int) – UV-grid oversampling factor. The PSF is gridded onto an oversampling * npx grid (finer UV cells, FoV scaled by oversampling) and the dirty beam is cropped back to npx x npx, discarding the borders where the KB correction amplifies noise/aliasing. Default is 2; set to 1 for the un-padded native-resolution behaviour.

Attributes:
grid_shape

Image grid shape (npx, npx).

grid_shape_os

Oversampled UV-grid shape (oversampling * npx,) * 2.

kernel_beta

KB shape parameter (2.34 * W, standard choice).

bw: float = 1000000000.0#
dec: float = 0.5585053606381855#
fov: float = 0.4#
freq: float = 2000000000.0#
property grid_shape: Tuple[int, int]#

Image grid shape (npx, npx).

property grid_shape_os: Tuple[int, int]#

Oversampled UV-grid shape (oversampling * npx,) * 2.

kernel_W: int = 7#
property kernel_beta: float#

KB shape parameter (2.34 * W, standard choice).

lat: float = 0.6108652381980153#
n_freqs: int = 5#
n_times: int = 5#
npx: int = 256#
oversampling: int = 2#
t_0: float = 1.0#
track_h: float = 0.5#
argosim.optim_utils.antenna_to_beam(antenna, config: ObservationConfig)[source]#

Antenna positions to normalised dirty beam.

Differentiable forward model mapping antenna ENU positions to the peak-normalised dirty beam (PSF) via baseline computation, UV-track generation, Kaiser-Bessel convolutional gridding, and inverse FFT.

Parameters:

  • antenna (jnp.ndarray) – Antenna ENU positions, shape (n_ant, 3) in metres.

  • config (ObservationConfig) – Static imaging / observation parameters.

Returns:

beam – Peak-normalised dirty beam, shape config.grid_shape.

Return type:

jnp.ndarray

argosim.optim_utils.beam_mse_loss(beam, target)[source]#

Beam MSE loss.

Normalised mean squared error between a (peak-normalised) dirty beam and a target image.

Parameters:

  • beam (jnp.ndarray) – Peak-normalised dirty beam.

  • target (jnp.ndarray) – Target image (e.g. a tight Gaussian).

Returns:

loss – Scalar normalised MSE.

Return type:

jnp.ndarray

argosim.optim_utils.centroid_penalty(antenna)[source]#

Centroid penalty.

Squared offset of the array centroid from the origin (E-N plane). Removes the loss’s translational degeneracy by pinning the centroid so the optimiser does not drift as a rigid body.

Parameters:

antenna (jnp.ndarray) – Antenna ENU positions, shape (n_ant, 3).

Returns:

penalty – Scalar non-negative penalty in m^2.

Return type:

jnp.ndarray

argosim.optim_utils.exp_decay_schedule(init_lr=25.0, end_lr=2.0, transition_steps=200, decay_rate=0.5)[source]#

Exp decay schedule.

Exponential learning-rate decay from init_lr toward end_lr.

Parameters:

  • init_lr (float) – Initial learning rate.

  • end_lr (float) – Asymptotic learning rate (lower bound).

  • transition_steps (int) – Number of steps over which to apply one decay factor.

  • decay_rate (float) – Multiplicative decay per transition_steps.

Returns:

schedule – Step-indexed learning-rate schedule.

Return type:

optax.Schedule

argosim.optim_utils.make_loss_fn(forward_fn, target_loss_fn, constraints: Sequence[Constraint])[source]#

Make loss fn.

Build a scalar loss combining a target / objective term and a set of weighted soft constraints.

Parameters:

  • forward_fn (Callable) – (antenna,) -> observable — differentiable. Typically a partially applied antenna_to_beam().

  • target_loss_fn (Callable) – (observable,) -> scalar target term.

  • constraints (sequence of Constraint) – Weighted soft constraints on antenna positions.

Returns:

loss_fn – (antenna,) -> scalar total loss.

Return type:

Callable

argosim.optim_utils.make_optimizer(name='adam', lr=1.0, **kwargs)[source]#

Make optimizer.

Build an optax optimiser by name.

Parameters:

  • name (str) – Optimiser name: "adam", "sgd", or "adamw".

  • lr (float or optax.Schedule) – Constant learning rate or an optax schedule.

  • **kwargs – Forwarded to the underlying optax constructor.

Returns:

optimizer – The configured optimiser.

Return type:

optax.GradientTransformation

argosim.optim_utils.optimise_array(antenna_init, loss_fn, optimizer, n_steps, snapshot_every: int | None = None, snapshot_steps: Sequence[int] | None = None, freeze_u: bool = True, verbose_every: int | None = 50)[source]#

Optimise array.

Run gradient-based optimisation of the antenna layout.

Parameters:

  • antenna_init (jnp.ndarray) – Initial antenna ENU positions, shape (n_ant, 3).

  • loss_fn (Callable) – (antenna,) -> scalar total loss, e.g. from make_loss_fn().

  • optimizer (optax.GradientTransformation) – The optimiser to step with, e.g. from make_optimizer().

  • n_steps (int) – Number of optimisation steps to run.

  • snapshot_every (int, optional) – If given, store antenna positions and loss every this many steps (uniform spacing).

  • snapshot_steps (sequence of int, optional) – Explicit step indices at which to store a snapshot. Useful for non-uniform spacing (e.g. log-spaced: dense early, sparse late). Overrides snapshot_every when provided. Step 0 (initial) and step n_steps (final) are always included.

  • freeze_u (bool) – If True, zero out the U (vertical) gradient component so the array stays flat.

  • verbose_every (int, optional) – Print a status line every this many steps. Set to None or 0 to silence.

Returns:

result – Keys: antenna (final positions, jnp.ndarray), loss_history (np.ndarray, length n_steps), snapshots_antenna (np.ndarray, shape (n_frames, n_ant, 3)), snapshots_loss (np.ndarray, length n_frames), snapshots_step (np.ndarray of int, length n_frames, the step index of each stored snapshot), snapshot_every (int or None).

Return type:

dict

argosim.optim_utils.radius_penalty(antenna, r_max)[source]#

Radius penalty.

Smooth ReLU barrier on each antenna’s distance from the array centroid above r_max. Models a circular site footprint.

Parameters:

  • antenna (jnp.ndarray) – Antenna ENU positions, shape (n_ant, 3).

  • r_max (float) – Maximum allowed radius from the centroid in metres.

Returns:

penalty – Scalar non-negative penalty.

Return type:

jnp.ndarray

argosim.optim_utils.spacing_penalty(antenna, d_min)[source]#

Spacing penalty.

Smooth ReLU barrier on antenna pairs closer than d_min in the E-N plane. Each violating pair contributes (1 - dist/d_min)^2.

Parameters:

  • antenna (jnp.ndarray) – Antenna ENU positions, shape (n_ant, 3).

  • d_min (float) – Minimum allowed pairwise distance in metres.

Returns:

penalty – Scalar non-negative penalty.

Return type:

jnp.ndarray