Point Source
Point-source microlensing utilities (JAX).
This module provides JAX-friendly implementations of core operations for point-source gravitational microlensing by up to three point-mass lenses.
Functions
lens_eq
: complex lens equation mapping image-plane coordinatesz
to source-plane coordinatesw
.lens_eq_det_jac
: determinant of the Jacobian of the lens mapping, used to compute magnification through|det J|^{-1}
.critical_and_caustic_curves
: critical curves in the image plane and their mapped caustics in the source plane._images_point_source
: image positions for given source position(s) by solving the corresponding complex polynomial._images_point_source_sequential
: helper that tracks images across a sequence of sources by reusing previous roots as initialization.mag_point_source
: total point-source magnification (sum over images).
All functions are compatible with jax.jit
and support batched evaluation
where appropriate. Complex numbers encode 2D coordinates using the x + i y
convention.
Coordinate conventions
Binary/triple lens configurations use the mid-point coordinate system for coefficient construction.
mag_point_source
shifts input source-plane coordinates to the center of mass when required and handles the inverse shift internally.
References
Bartolic, F.
caustics
(MIT License) — original inspiration and some coefficient-generation routines, modified here for JAX and extended triple lens support.
- microjax.point_source.critical_and_caustic_curves(npts: int = 200, nlenses: int = 2, **params)
Compute critical curves and mapped caustics.
- Parameters
npts (int, optional) – Number of sampling points on the unit circle used to construct the critical-curve polynomial. Defaults to 200.
nlenses (int, optional) – Number of point-mass lenses (1, 2, or 3). Defaults to 2.
params (dict) –
Lens configuration parameters:
nlenses = 1
: no additional parameters.nlenses = 2
:s
(separation) andq
(mass ratiom2/m1
). Internally we seta = s / 2
ande1 = q / (1 + q)
.nlenses = 3
:s
,q
,q3
(third mass ratio),r3
(radius),psi
(angle). Internallya = s / 2
,e1 = q / (1 + q + q3)
,e2 = 1 / (1 + q + q3)
.
- Returns
(z_cr, z_ca)
wherez_cr
contains critical curves andz_ca
the mapped caustics. Both arrays have shape(N_branches, npts)
.- Return type
tuple[jax.Array, jax.Array]
Notes
For
nlenses = 1
the critical curve is the unit circle and the caustic collapses to the origin.Output is shifted from mid-point to center of mass for consistency with the rest of the library.
- microjax.point_source.lens_eq(z: Array, nlenses: int = 2, **params) Array
Lens equation mapping image-plane
z
to source-planew
.- Parameters
z (jax.Array) – Complex scalar or array of image-plane coordinates.
nlenses (int, optional) – Number of point-mass lenses (1, 2, or 3). Defaults to 2.
params (dict) – Lens configuration parameters. For
nlenses = 2
expecta
(half separation) ande1
(mass fraction at+a
). Fornlenses = 3
expecta
,r3
,psi
,e1
ande2
.
- Returns
Source-plane coordinates with the same shape as
z
.- Return type
jax.Array
Notes
All arithmetic is performed in complex form; gradients propagate through JAX as expected.
- microjax.point_source.lens_eq_det_jac(z: Array, nlenses: int = 2, **params) Array
Determinant of the Jacobian of the lens mapping at
z
.- Parameters
z (jax.Array) – Complex scalar or array of image-plane coordinates.
nlenses (int, optional) – Number of point-mass lenses (1, 2, or 3). Defaults to 2.
params (dict) – Lens configuration parameters matching those accepted by
lens_eq()
.
- Returns
Real array with the same shape as
z
storingdet J(z)
.- Return type
jax.Array
Notes
Point-source magnification is
|det J|^{-1}
for each image.
- microjax.point_source.mag_point_source(w, nlenses=2, **params)
Total point-source magnification for 1–3 lens configurations.
- Parameters
w (jax.Array) – Complex scalar or array of source-plane coordinates.
nlenses (int, optional) – Number of point-mass lenses (1, 2, or 3). Defaults to 2.
params (dict) –
Lens parameters depend on
nlenses
:nlenses = 1
: no additional parameters required.nlenses = 2
:s
(separation) andq
(mass ratiom2/m1
). Internallya = s / 2
ande1 = q / (1 + q)
; the source is shifted to the center of mass for the polynomial construction.nlenses = 3
:s
,q
,q3
,r3
andpsi
. Internallya = s / 2
,e1 = q / (1 + q + q3)
,e2 = 1 / (1 + q + q3)
, and the same center-of-mass shift is applied.
- Returns
Real-valued magnification with the same shape as
w
.- Return type
jax.Array
Notes
Magnification is computed as the sum of
|det J|^{-1}
over valid image branches returned by_images_point_source
.