From 18c4fabd6b2f9f6fc6b0dd74c02ab14d67ec4bef Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 10:51:14 -0500
Subject: [PATCH 01/10] Enforce that the transit chord lies below the stellar
 equator in Star class

---
 fleck/core.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/fleck/core.py b/fleck/core.py
index 1482229..eff5dd1 100644
--- a/fleck/core.py
+++ b/fleck/core.py
@@ -278,6 +278,19 @@ def light_curve(self, spot_lons, spot_lats, spot_radii, inc_stellar,
                             np.cos(Omega) * np.sin(omega + f) * np.cos(I))
             Z = planet.a * np.sin(omega + f) * np.sin(I)
 
+            # Enforce that the transit chord lies below the stellar equator
+            # in the same "observer oriented" frame used for the spots/plot.
+            in_front = (X < 0) & (np.abs(Y) < 1 + planet.rp)
+            if np.any(in_front):
+                idx = np.where(in_front)[0]
+                # approximate mid-transit as the in-front point with minimum |Y|
+                mid_idx = idx[np.argmin(np.abs(Y[in_front]))]
+                z_star_mid = -Z[mid_idx]
+                if z_star_mid > 0:
+                    # Planet is crossing the lat > 0 hemisphere; reflect it
+                    # across the equator so it crosses lat < 0 instead.
+                    Z = -Z
+
             # Create a shapely circle object for the planet's silhouette only
             # when the planet is in front of the star, otherwise append `None`
             planet_disk = [circle([-Y[i], -Z[i]], planet.rp)

From 253e5f1e18fcc7af23693790d4b7a7b62fc4b82f Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:01:30 -0500
Subject: [PATCH 02/10] Add settings.json to .gitignore

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index 22ce1b4..7f8f5a0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -55,6 +55,7 @@ distribute-*.tar.gz
 .project
 .pydevproject
 .settings
+settings.json
 
 docs/.DS_Store
 .DS_Store

From 07c5d207a95182b6f8f95fea1af51f0a7b3fa68c Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:06:37 -0500
Subject: [PATCH 03/10] Changing to using latitudes instead of colatitudes for
 fleck.jax spots

---
 fleck/jax.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index aa619d9..e358d31 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -55,7 +55,7 @@ def __init__(
         lon : array
             Active region longitudes in radians on (0, 2pi)
         lat : array
-            Active region latitudes in radians on (0, pi)
+            Active region latitudes in radians on (-pi/2, pi/2)
         rad : array
             Active region radii in units of stellar radii
         spectrum : array
@@ -203,7 +203,7 @@ def spot_coords(self, times=None, t0_rot=0):
 
         """
         Limits:
-        lat: (0, pi)
+        lat: (-pi/2, pi/2)
         lon: (0, 2pi)
         rad: (0, None)
         contrast: (0, inf)
@@ -220,7 +220,6 @@ def spot_coords(self, times=None, t0_rot=0):
         lon = jnp.expand_dims(self.lon, [0, 2, 3])
         lat = jnp.expand_dims(self.lat, [0, 2, 3])
         rad = jnp.expand_dims(self.rad, [0, 2, 3])
-        contrast = jnp.expand_dims(contrast, [0, 3])
         inclination = jnp.expand_dims(jnp.asarray(self.inclination), [0, 1, 2])
 
         comp_inclination = np.pi / 2 - inclination
@@ -232,13 +231,13 @@ def spot_coords(self, times=None, t0_rot=0):
         cos_c_inc = jnp.cos(comp_inclination)
 
         spot_position_x = (
-            jnp.cos(phi - np.pi / 2) * sin_c_inc * sin_lat +
-            cos_c_inc * cos_lat
+            jnp.cos(phi - np.pi / 2) * sin_c_inc * cos_lat +
+            cos_c_inc * sin_lat
         )
-        spot_position_y = -jnp.sin(phi - np.pi / 2) * sin_lat
+        spot_position_y = -jnp.sin(phi - np.pi / 2) * cos_lat
         spot_position_z = (
-            cos_lat * sin_c_inc -
-            jnp.sin(phi) * cos_c_inc * sin_lat
+            sin_lat * sin_c_inc -
+            jnp.sin(phi) * cos_c_inc * cos_lat
         )
 
         rsq = spot_position_x ** 2 + spot_position_y ** 2
@@ -261,7 +260,7 @@ def add_spot(self, lon, lat, rad, contrast=None, temperature=None, spectrum=None
         lon : float
             Active region longitudes in radians on (0, 2pi)
         lat : float
-            Active region latitudes in radians on (0, pi)
+            Active region latitudes in radians on (-pi/2, pi/2)
         rad : float
             Active region radii in units of stellar radii
         contrast : float

From 4f72c50346e7dbf6add893ab8cb0aec4e4fc6892 Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:11:54 -0500
Subject: [PATCH 04/10] Minor formatting changes

---
 fleck/jax.py | 179 ++++++++++++++++++++++++++++++++-------------------
 1 file changed, 111 insertions(+), 68 deletions(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index e358d31..3ce6a12 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -30,7 +30,8 @@ class ActiveStar:
     optional planetary transit models and spot occultations.
     """
 
-    n_mc = 10_000  # Number of Monte Carlo samples to use when computing planet+spot overlap
+    # Number of Monte Carlo samples to use when computing planet+spot overlap
+    n_mc = 10_000
     key = random.PRNGKey(0)  # random key seed
 
     def __init__(
@@ -137,7 +138,7 @@ def rotation_model(self, f0=0, t0_rot=0, u1=0, u2=0):
         radial_coord = 1 - jnp.geomspace(1e-5, 1, 100)[::-1]
         unspotted_total_flux = trapezoid(
             y=(
-                2 * np.pi * radial_coord *
+                2 * jnp.pi * radial_coord *
                 self.limb_darkening(radial_coord, u1, u2)
             ),
             x=radial_coord
@@ -145,7 +146,7 @@ def rotation_model(self, f0=0, t0_rot=0, u1=0, u2=0):
 
         # Morris 2020 Eqn 6-7
         spot_model = f0 - jnp.sum(
-            np.pi * rad ** 2 *
+            jnp.pi * rad ** 2 *
             (1 - contrast) *
             self.limb_darkening(mu, u1, u2) *
             mask_behind_star,
@@ -171,19 +172,23 @@ def spot_coords(self, times=None, t0_rot=0):
         Returns
         -------
         spot_position_x : array
-            x-position of the active region in the observer oriented coordinate system [1]_.
+            x-position of the active region in the observer oriented
+            coordinate system [1]_.
         spot_position_y : array
-            y-position of the active region in the observer oriented coordinate system [1]_.
+            y-position of the active region in the observer oriented
+            coordinate system [1]_.
         spot_position_z : array
-            y-position of the active region in the observer oriented coordinate system [1]_.
+            y-position of the active region in the observer oriented
+            coordinate system [1]_.
         major_axis : array
-            Apparent semimajor axis of the circular active region, which is elliptical when
-            projected active onto the sky plane (in general)
+            Apparent semimajor axis of the circular active region, which is
+            elliptical when projected active onto the sky plane (in general)
         minor_axis : array
-            Apparent semiminor axis of the circular active region, which is elliptical when
-            projected active onto the sky plane (in general)
+            Apparent semiminor axis of the circular active region, which is
+            elliptical when projected active onto the sky plane (in general)
         angle : array
-            Angle between the +x-axis and the projected active region's semimajor axis
+            Angle between the +x-axis and the projected active region's
+            semimajor axis
         rad : array
             Active region radius [stellar radii]
         contrast: array
@@ -216,14 +221,15 @@ def spot_coords(self, times=None, t0_rot=0):
         3. inclination
         """
 
-        phase = jnp.expand_dims(2 * np.pi * (times - t0_rot) / self.P_rot, [1, 2, 3])
+        phase = jnp.expand_dims(2 * jnp.pi * (times - t0_rot) / self.P_rot,
+                                [1, 2, 3])
         lon = jnp.expand_dims(self.lon, [0, 2, 3])
         lat = jnp.expand_dims(self.lat, [0, 2, 3])
         rad = jnp.expand_dims(self.rad, [0, 2, 3])
         inclination = jnp.expand_dims(jnp.asarray(self.inclination), [0, 1, 2])
 
-        comp_inclination = np.pi / 2 - inclination
-        phi = np.pi / 2 - phase - lon
+        comp_inclination = jnp.pi / 2 - inclination
+        phi = jnp.pi / 2 - phase - lon
 
         sin_lat = jnp.sin(lat)
         cos_lat = jnp.cos(lat)
@@ -231,10 +237,10 @@ def spot_coords(self, times=None, t0_rot=0):
         cos_c_inc = jnp.cos(comp_inclination)
 
         spot_position_x = (
-            jnp.cos(phi - np.pi / 2) * sin_c_inc * cos_lat +
+            jnp.cos(phi - jnp.pi / 2) * sin_c_inc * cos_lat +
             cos_c_inc * sin_lat
         )
-        spot_position_y = -jnp.sin(phi - np.pi / 2) * cos_lat
+        spot_position_y = -jnp.sin(phi - jnp.pi / 2) * cos_lat
         spot_position_z = (
             sin_lat * sin_c_inc -
             jnp.sin(phi) * cos_c_inc * cos_lat
@@ -251,7 +257,8 @@ def spot_coords(self, times=None, t0_rot=0):
             major_axis, minor_axis, angle, rad, contrast
         )
 
-    def add_spot(self, lon, lat, rad, contrast=None, temperature=None, spectrum=None):
+    def add_spot(self, lon, lat, rad, contrast=None, temperature=None,
+                 spectrum=None):
         """
         Add an active region to the stellar model.
 
@@ -263,12 +270,16 @@ def add_spot(self, lon, lat, rad, contrast=None, temperature=None, spectrum=None
             Active region latitudes in radians on (-pi/2, pi/2)
         rad : float
             Active region radii in units of stellar radii
-        contrast : float
+        contrast : float or array
             Ratio of the active region's flux to the photospheric
-            flux at each ``ActiveStar.wavelength``
-        spectrum : float
+            flux at each ``ActiveStar.wavelength``.
+        temperature : float
+            Effective temperature of the active region. If provided
+            (and neither ``contrast`` nor ``spectrum`` are given),
+            a blackbody spectrum will be computed.
+        spectrum : float or array
             The spectrum of the active region on the same wavelength
-            grid is ``ActiveStar.phot``
+            grid as ``ActiveStar.phot``.
         """
         if contrast is None and spectrum is None and temperature is not None:
             self.phot = self._blackbody(self.wavelength, self.T_eff)
@@ -310,15 +321,15 @@ def limb_darkening(self, mu, u1, u2):
         Compute quadratic limb darkening as a function of :math:`\\mu`.
         """
         return (
-            1 / np.pi *
+            1 / jnp.pi *
             (1 - u1 * (1 - mu) - u2 * (1 - mu) ** 2) /
             (1 - u1 / 3 - u2 / 6)
         )
 
     @jit
     def transit_model(self, t0, period, rp, a, inclination,
-                      omega=np.pi / 2, ecc=0, f0=1, t0_rot=0,
-                      u1=0, u2=0):
+                      omega=jnp.pi/2, ecc=0., f0=1., t0_rot=0.,
+                      u1=0., u2=0.):
         """
         Compute spectrophotometry with rotation and a planetary transit.
 
@@ -331,8 +342,8 @@ def transit_model(self, t0, period, rp, a, inclination,
             Mid-transit time
         period : float
             Orbital period of the transiting planet
-        rp : float
-            Exoplanet radius in units of stellar radii
+        rp : float or array
+            Exoplanet radius in units of stellar radii for each wavelength
         a : float
             Planetary semi-major axis in units of stellar radii
         inclination : float
@@ -345,10 +356,10 @@ def transit_model(self, t0, period, rp, a, inclination,
             Out-of-transit flux for an unspotted star, default is one.
         t0_rot : float
             Zero-point in time for stellar rotation, default is zero
-        u1 : float
-            Limb-darkening parameter :math:`u_1`
-        u2 : float
-            Limb-darkening parameter :math:`u_2`
+        u1 : float or array
+            Limb-darkening parameter :math:`u_1` for each wavelength
+        u2 : float or array
+            Limb-darkening parameter :math:`u_2` for each wavelength
 
         Returns
         -------
@@ -358,9 +369,11 @@ def transit_model(self, t0, period, rp, a, inclination,
             The apparent squared ratio of planet-to-star radius with stellar
             spectral contamination by active regions
          X : array
-            x-position of the planet in the observer oriented coordinate system [1]_.
+            x-position of the planet in the observer oriented coordinate
+            system [1]_.
          Y : array
-            y-position of the planet in the observer oriented coordinate system [1]_.
+            y-position of the planet in the observer oriented coordinate
+            system [1]_.
 
         References
         ----------
@@ -377,16 +390,16 @@ def transit_model(self, t0, period, rp, a, inclination,
         ) = self.spot_coords(t0_rot=t0_rot)
 
         rsq = spot_position_x ** 2 + spot_position_y ** 2
-        mu = jnp.sqrt(1 - rsq)
+        mu = jnp.sqrt(1. - rsq)
         mask_behind_star = jnp.where(
-            spot_position_z < 0, mu, 0
+            spot_position_z < 0., mu, 0.
         )
 
-        radial_coord = 1 - jnp.geomspace(1e-5, 1, 100)[::-1]
+        radial_coord = 1. - jnp.geomspace(1e-5, 1., 100)[::-1]
 
         unspotted_total_flux = trapezoid(
             y=(
-                2 * np.pi * radial_coord[:, None] *
+                2. * jnp.pi * radial_coord[:, None] *
                 self.limb_darkening(
                     radial_coord[:, None], *u_ld.T
                 )
@@ -402,15 +415,15 @@ def transit_model(self, t0, period, rp, a, inclination,
 
         # Morris 2020 Eqn 6-7
         out_of_transit = f0 - jnp.sum(
-            np.pi * rad ** 2 *
-            (1 - contrast) *
+            jnp.pi * rad ** 2 *
+            (1. - contrast) *
             limb_dark *
             mask_behind_star /
             unspotted_total_flux[None, None, :, None],
             axis=1
         )
 
-        f_S = rad ** 2 * mu * (spot_position_z < 0).astype(int)
+        f_S = rad ** 2 * mu * (spot_position_z < 0.).astype(int)
 
         # compute the transit model
         mean_anomaly = 2 * np.pi * (self.times - t0) / period
@@ -418,10 +431,10 @@ def transit_model(self, t0, period, rp, a, inclination,
             *jaxoplanet.core.kepler(M=mean_anomaly, ecc=ecc)
         )
 
-        # Winn 2011 Eqn 1
-        r = a * (1 - ecc ** 2) / (1 + ecc * jnp.cos(true_anomaly))
+        # Winn 2011 Eqn 1: sky-projected star-planet distance
+        r = a * (1. - ecc ** 2) / (1. + ecc * jnp.cos(true_anomaly))
 
-        # Winn 2011 Eqn 3-4
+        # Winn 2011 Eqn 3-4: sky-projected coordinates
         X = -r * jnp.cos(omega + true_anomaly)
         Y = -r * jnp.sin(omega + true_anomaly) * jnp.cos(inclination)
 
@@ -463,13 +476,14 @@ def transit_model(self, t0, period, rp, a, inclination,
         depth_ratio = contaminated_max_depth / uncontaminated_max_depth
         apparent_rprs2 = rp ** 2 * depth_ratio
 
+        # Planet–spot distance in the (X, Y) plane
         planet_spot_distance = jnp.hypot(
             spot_position_y - X[:, None, None, None],
-            spot_position_x - Y[:, None, None, None]
+            spot_position_x - Y[:, None, None, None],
         )
         occultation_possible = jnp.squeeze(
             (planet_spot_distance < (major_axis + rp.mean())) &
-            (spot_position_z < 0)
+            (spot_position_z < 0.)
         )
 
         @jit
@@ -492,27 +506,30 @@ def time_step(
                     radius=rp,
                     occultation_possible=occultation_possible[j],
                 ),
-                lambda *args: jnp.zeros((spot_position_x.shape[1], self.n_mc), dtype=bool),
+                lambda *args: jnp.zeros((spot_position_x.shape[1], self.n_mc),
+                                        dtype=bool),
             )
 
         occultation_per_time_per_spot_per_mc_sample = lax.scan(
             time_step, 0.0, jnp.arange(self.times.shape[0])
-        )[1]  # shape: (n_times, n_spots, n_mc_samples)
+        )[1]
+        # (n_times, n_spots, n_mc_samples)
 
         frac_occulted_per_time_per_spot = jnp.count_nonzero(
             occultation_per_time_per_spot_per_mc_sample, axis=2
         ) / self.n_mc
 
         occultation = (
-            (1 - contrast) *
+            (1. - contrast) *
             jnp.expand_dims(frac_occulted_per_time_per_spot, axis=(2, 3))
         )
-        scaled_occultation = (1 - contaminated_transit) * jnp.sum(occultation, axis=1)[..., 0]
+        scaled_occultation = ((1. - contaminated_transit)
+                              * jnp.sum(occultation, axis=1)[..., 0])
 
         spectrum_at_transit = time_series_spectrum[t_ind]
 
         return (
-            out_of_transit[..., 0] * (contaminated_transit + scaled_occultation),
+            out_of_transit[..., 0]*(contaminated_transit+scaled_occultation),
             apparent_rprs2, X, Y,
             spectrum_at_transit
         )
@@ -524,9 +541,11 @@ def _area_union_per_time(
     ):
         # Monte Carlo sampling for points inside the planet's disk:
         key, subkey = random.split(self.key)
-        theta_p = random.uniform(key, minval=0, maxval=2 * np.pi, shape=(self.n_mc,))
+        theta_p = random.uniform(key, minval=0, maxval=2 * jnp.pi,
+                                 shape=(self.n_mc,))
         key, subkey = random.split(key)
-        rad_p = radius * random.uniform(subkey, minval=0, maxval=1, shape=(self.n_mc,)) ** 0.5
+        rad_p = radius * random.uniform(subkey, minval=0, maxval=1,
+                                        shape=(self.n_mc,)) ** 0.5
         xp = rad_p * jnp.cos(theta_p) + x0_circle
         yp = rad_p * jnp.sin(theta_p) + y0_circle
 
@@ -542,12 +561,15 @@ def _area_union_per_time(
 
         @jit
         def find_overlap(k):
-            # find overlap between the planet and the elliptical region (projected circular spot)
+            # find overlap between the planet and the elliptical region
+            # (projected circular spot)
             in_ellipse = jnp.hypot(
                 ((xp - x0_ellipse[k]) * jnp.cos(jnp.radians(angle[k])) +
-                 (yp - y0_ellipse[k]) * jnp.sin(jnp.radians(angle[k]))) / alpha[k],
+                 (yp - y0_ellipse[k]) * jnp.sin(jnp.radians(angle[k]))
+                 ) / alpha[k],
                 ((xp - x0_ellipse[k]) * jnp.sin(jnp.radians(angle[k])) -
-                 (yp - y0_ellipse[k]) * jnp.cos(jnp.radians(angle[k]))) / beta[k]
+                 (yp - y0_ellipse[k]) * jnp.cos(jnp.radians(angle[k]))
+                 ) / beta[k]
             ) < 1
 
             return in_ellipse & on_star
@@ -563,7 +585,8 @@ def spot_step(carry, k, occultation_possible=occultation_possible):
                 lambda *args: jnp.zeros(self.n_mc, dtype=bool)
             )
 
-        monte_carlo_occulted_area = lax.scan(spot_step, 0, jnp.arange(x0_ellipse.shape[0]))[1]
+        monte_carlo_occulted_area = lax.scan(
+            spot_step, 0, jnp.arange(x0_ellipse.shape[0]))[1]
 
         return monte_carlo_occulted_area
 
@@ -588,8 +611,8 @@ def plot_star(self, t0, rp, a, inclination,
         t0_rot : float
             Zero-point in time for stellar rotation, default is zero
         multiply_radii : float
-            Visually represent scaled-up active regions where the radii are increased
-            by factor ``multiply_radii``, default is one.
+            Visually represent scaled-up active regions where the radii are
+            increased by factor ``multiply_radii``, default is one.
         ax : matplotlib.axes.Axes
             Add the visualization to this matplotlib axis
         annotate : bool
@@ -701,20 +724,27 @@ def bin_spectrum(spectrum, bins=None, log=True, min=None, max=None, **kwargs):
     Parameters
     ----------
     spectrum : `specutils.Spectrum1D`
+        Spectrum to be binned
+    bins : int or ~numpy.ndarray
+        Number of bins, or the bin edges
     log : bool
         If true, compute bin edges based on the log base 10 of
         the frequency.
-    bins : int or ~numpy.ndarray
-        Number of bins, or the bin edges
+    min : ~astropy.units.Quantity
+        Minimum wavelength to include in the binned spectrum.
+    max : ~astropy.units.Quantity
+        Maximum wavelength to include in the binned spectrum.
+    **kwargs : dict
+        Additional keyword arguments passed to `scipy.stats.binned_statistic`
 
     Returns
     -------
-    new_spectrum :
+    new_spectrum : `specutils.Spectrum1D`
+        Binned spectrum
     """
-    nirspec_wl_range = (spectrum.wavelength > min) & (spectrum.wavelength < max)
-
-    wavelength = spectrum.wavelength[nirspec_wl_range]
-    flux = spectrum.flux[nirspec_wl_range]
+    in_range = (spectrum.wavelength > min) & (spectrum.wavelength < max)
+    wavelength = spectrum.wavelength[in_range]
+    flux = spectrum.flux[in_range]
 
     if log:
         wl_axis = np.log10(wavelength.to(u.um).value)
@@ -740,20 +770,33 @@ def bin_spectrum(spectrum, bins=None, log=True, min=None, max=None, **kwargs):
     nans = np.isnan(bs.statistic)
     interp_fluxes = bs.statistic.copy()
     if np.any(nans) and all(
-        map(lambda x: len(x) > 0, [wl_bins[nans], wl_bins[~nans], bs.statistic[~nans]])
+        map(lambda x: len(x) > 0, [wl_bins[nans], wl_bins[~nans],
+                                   bs.statistic[~nans]])
     ):
-        interp_fluxes[nans] = np.interp(wl_bins[nans], wl_bins[~nans], bs.statistic[~nans])
+        interp_fluxes[nans] = np.interp(wl_bins[nans], wl_bins[~nans],
+                                        bs.statistic[~nans])
     return Spectrum1D(
-        flux=interp_fluxes * flux.unit, spectral_axis=wl_bins, meta=spectrum.meta
+        flux=interp_fluxes * flux.unit, spectral_axis=wl_bins,
+        meta=spectrum.meta
     )
 
 
 def spectral_binning(y, all_x, all_y):
     """
     Spectral binning via trapezoidal approximation.
+
+    Parameters
+    ----------
+    y : array
+        Values to be binned
+    all_x : array
+        Full x-axis of the spectrum
+    all_y : array
+        Full y-axis of the spectrum
     """
     min_ind = np.argwhere(all_y == y[0])[0, 0]
     max_ind = np.argwhere(all_y == y[-1])[0, 0]
     if max_ind > min_ind and y.shape == all_x[min_ind:max_ind + 1].shape:
-        return np.trapz(y, all_x[min_ind:max_ind + 1]) / (all_x[max_ind] - all_x[min_ind])
+        return (np.trapz(y, all_x[min_ind:max_ind + 1])
+                / (all_x[max_ind] - all_x[min_ind]))
     return np.nan

From 1cf96f94541fe38b6ce2c80aebdd42b02af92b52 Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:22:23 -0500
Subject: [PATCH 05/10] Add contrast parameter to ActiveStar class init and
 update related methods

---
 fleck/jax.py | 312 ++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 274 insertions(+), 38 deletions(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index 3ce6a12..a035783 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -46,7 +46,8 @@ def __init__(
         inclination=empty,
         wavelength=None,
         phot=None,
-        P_rot=3.3
+        P_rot=3.3,
+        contrast=empty,
     ):
         """
         Parameters
@@ -73,6 +74,11 @@ def __init__(
             Photospheric flux at each ``wavelength``.
         P_rot : float
             Stellar rotation period
+        contrast : float or array
+            Ratio of the active region flux to the photospheric
+            flux at each ``wavelength``. If provided, this will be
+            used directly; otherwise it is computed from ``spectrum``
+            and ``phot`` when possible.
         """
         self.times = jnp.array(times)
         self.lon = jnp.array(lon)
@@ -82,10 +88,104 @@ def __init__(
         self.T_eff = T_eff
         self.temperature = jnp.array(temperature)
         self.inclination = inclination
-        self.wavelength = wavelength
-        self.phot = phot
+        # Allow wavelength/phot to remain optional, but store them as
+        # JAX arrays when provided.
+        self.wavelength = None if wavelength is None else jnp.array(wavelength)
+        self.phot = None if phot is None else jnp.array(phot)
         self.P_rot = P_rot
 
+        # Raw contrast as passed by the user (may be scalar, 1D, 2D, or
+        # empty). We normalise shapes after dimension inference.
+        if contrast is None or getattr(contrast, "size", 0) == 0:
+            self.contrast = jnp.array([])
+        else:
+            self.contrast = jnp.array(contrast)
+
+        # ------------------------------------------------------------------
+        # Dimension bookkeeping: infer n_times, n_spots, n_lambda
+        # ------------------------------------------------------------------
+        self.n_times = int(self.times.size)
+        self.n_spots = int(self.rad.size)
+        self.n_lambda = None
+
+        if (
+            self.wavelength is not None and
+            getattr(self.wavelength, "size", 0) != 0
+        ):
+            # 1) Prefer the explicit wavelength grid, if provided
+            self.n_lambda = int(self.wavelength.size)
+
+        elif getattr(self.spectrum, "size", 0) != 0:
+            # 2) Otherwise, use the spectrum shape
+
+            if self.spectrum.ndim == 1:
+                # Treat a 1D spectrum as (n_lambda,)
+                self.n_lambda = int(self.spectrum.shape[0])
+            else:
+                # Assume the last axis is wavelength
+                self.n_lambda = int(self.spectrum.shape[-1])
+
+        elif getattr(self.phot, "size", 0) != 0:
+            # 3) Otherwise, use the photosphere
+
+            self.n_lambda = self.phot.size
+
+        elif getattr(self.contrast, "size", 0) != 0:
+            # 4) Finally, fall back to the contrast shape, using n_spots to
+            # break degeneracies between (n_spots,) and (n_lambda,)
+
+            c = self.contrast
+            if c.ndim == 0:
+                self.n_lambda = 1
+            elif c.ndim == 1:
+                if self.n_spots > 0 and c.shape[0] == self.n_spots:
+                    # One contrast per spot, same for all wavelengths
+                    self.n_lambda = 1
+                else:
+                    # One contrast per wavelength bin, same for all spots
+                    self.n_lambda = int(c.shape[0])
+            elif c.ndim == 2:
+                # Last axis is wavelength
+                self.n_lambda = int(c.shape[-1])
+
+        else:
+            # Absolute last resort: a single "white-light" wavelength bin
+            self.n_lambda = 1
+
+        # Normalise the supplied contrast argument into something at least
+        # 2D where possible.
+        if self.contrast.size != 0:
+            if self.contrast.ndim == 0:
+                # Single scalar contrast: broadcast to all spots in a single
+                # wavelength bin.
+                self.contrast = self.contrast * jnp.ones((self.n_spots, 1))
+            elif self.contrast.ndim == 1:
+                if self.n_spots > 0 and self.contrast.shape[0] == self.n_spots:
+                    # Per-spot contrast for a single wavelength bin.
+                    self.contrast = self.contrast[:, None]
+                else:
+                    # Per-wavelength contrast for a single spot.
+                    self.contrast = self.contrast[None, :]
+            # If contrast is already 2D+ we assume the user passed
+            # (n_spots, n_lambda) and leave it as-is.
+
+        # If we have no phot, set it to unity
+        self.phot = (jnp.asarray(self.phot) if self.phot is not None
+                     else jnp.ones((self.n_lambda,)))
+
+        # If an explicit contrast was not provided but we do have a
+        # spectrum and photosphere, derive the contrast now.
+        if self.contrast.size == 0 and self.spectrum.size != 0:
+            derived_contrast = self.spectrum / self.phot[None, :]
+            if derived_contrast.ndim == 1:
+                derived_contrast = derived_contrast[None, :]
+            self.contrast = derived_contrast
+
+        # If we have a contrast and a photosphere but no spectrum,
+        # construct a spectrum that is consistent with both.
+        if self.spectrum.size == 0 and self.contrast.size != 0:
+            self.spectrum = self.contrast * self.phot[None, :]
+
     def tree_flatten(self):
         children = (
             self.times,
@@ -99,6 +199,7 @@ def tree_flatten(self):
             self.wavelength,
             self.phot,
             self.P_rot,
+            self.contrast,
         )
         aux_data = None
         return children, aux_data
@@ -198,14 +299,50 @@ def spot_coords(self, times=None, t0_rot=0):
         ----------
         .. [1]  Fabrycky & Winn (2009) https://arxiv.org/abs/0902.0737
         """
-        contrast = self.spectrum / self.phot[None, :]
-
-        if contrast.ndim == 1:
-            contrast = contrast[None, :]
-
         if times is None:
             times = self.times
 
+        # Use the dimensions inferred in __init__
+        n_spots = int(self.n_spots)
+        n_lambda = int(self.n_lambda)
+
+        # ------------------------------------------------------------------
+        # Build a per-spot, per-wavelength contrast array of shape
+        # (n_spots, n_lambda), then broadcast to (1, n_spots, n_lambda, 1)
+        # for use with the f_S tensor in transit_model.
+        # ------------------------------------------------------------------
+        if self.contrast.size != 0:
+            base_contrast = jnp.asarray(self.contrast)
+        elif self.spectrum.size != 0:
+            base_contrast = jnp.asarray(self.spectrum) / self.phot[None, :]
+        else:
+            # No contrast information available.
+            # Default to maximum contrast: spot == 0
+            base_contrast = jnp.zeros((n_spots, n_lambda))
+
+        # Normalize to exactly (n_spots, n_lambda)
+        if n_spots == 0:
+            # No spots: keep a well-defined empty (0, n_lambda) array
+            # to avoid ambiguous reshapes.
+            base_contrast = jnp.zeros((0, n_lambda))
+        elif base_contrast.ndim == 0:
+            base_contrast = jnp.full((n_spots, n_lambda), base_contrast)
+        elif base_contrast.ndim == 1:
+            if base_contrast.shape[0] == n_spots:
+                # One contrast per spot, broadcast over wavelength.
+                base_contrast = base_contrast[:, None]*jnp.ones((1, n_lambda))
+            else:
+                # One contrast per wavelength, broadcast over spots.
+                base_contrast = base_contrast[None, :]*jnp.ones((n_spots, 1))
+        elif base_contrast.shape[0] == 1 and n_spots > 1:
+            # Single-spot input, broadcast along spot axis first.
+            base_contrast = jnp.broadcast_to(
+                base_contrast, (n_spots,)+base_contrast.shape[1:])
+
+        # Final contrast tensor used downstream:
+        # shape: (1, n_spots, n_lambda, 1)
+        contrast = base_contrast[None, :, :, None]
+
         """
         Limits:
         lat: (-pi/2, pi/2)
@@ -281,26 +418,58 @@ def add_spot(self, lon, lat, rad, contrast=None, temperature=None,
             The spectrum of the active region on the same wavelength
             grid as ``ActiveStar.phot``.
         """
-        if contrast is None and spectrum is None and temperature is not None:
-            self.phot = self._blackbody(self.wavelength, self.T_eff)
+        if contrast is not None:
+            # User-supplied contrast: convert to spectrum
+            contrast = jnp.array(contrast)
+            if contrast.ndim == 0:
+                # Scalar contrast: broadcast to all wavelengths
+                contrast = contrast[None] * jnp.ones(self.n_lambda)
+            elif contrast.ndim == 1:
+                # 1D contrast are defined on the wavelength grid
+                contrast = contrast[None, :]
+            # Derive the spectrum from the contrast
+            spectrum = contrast * self.phot[None, :]
+        elif spectrum is not None:
+            # User-supplied spectrum: ensure correct shape
+            spectrum = jnp.array(spectrum)
+            if spectrum.ndim == 0:
+                # Scalar spectrum: broadcast to all wavelengths
+                spectrum = spectrum[None] * jnp.ones(self.n_lambda)
+            elif spectrum.ndim == 1:
+                # 1D spectra are defined on the wavelength grid
+                spectrum = spectrum[None, :]
+            # Derive the contrast from the spectrum
+            contrast = spectrum / self.phot[None, :]
+        elif spectrum is None and temperature is not None:
+            # No contrast or spectrum, but a temperature: compute a blackbody
             spectrum = self._blackbody(self.wavelength, temperature)
+            contrast = spectrum / self.phot[None, :]
+            if contrast.ndim == 1:
+                contrast = contrast[None, :]
 
-        for attr, new_value in zip("lon, lat, rad, spectrum, temperature".split(', '),
-                                   [lon, lat, rad, spectrum, temperature]):
+        for attr, new_value in zip(
+            "lon, lat, rad, spectrum, temperature, contrast".split(', '),
+            [lon, lat, rad, spectrum, temperature, contrast],
+        ):
 
             prop = getattr(self, attr)
 
             if not hasattr(new_value, 'ndim'):
                 new_value = jnp.array([new_value])
 
-            if prop is not None:
-                if prop.ndim > 1 or (len(prop) > 1 and len(prop) == len(new_value)):
+            if prop is not None and getattr(prop, "size", 0) != 0:
+                if prop.ndim > 1 or (
+                    len(prop) > 1 and len(prop) == len(new_value)
+                ):
                     new_value = jnp.vstack([prop, new_value])
                 else:
                     new_value = jnp.concatenate([prop, new_value])
 
                 setattr(self, attr, new_value)
 
+        # Update n_spots
+        self.n_spots += 1
+
     @jit
     def _blackbody(self, wavelength_meters, temperature):
         """
@@ -379,6 +548,15 @@ def transit_model(self, t0, period, rp, a, inclination,
         ----------
         .. [1] Fabrycky & Winn (2009) https://arxiv.org/abs/0902.0737
         """
+        # Photosphere: ensure a 1D array of length n_lambda
+        n_lambda = int(self.n_lambda)
+        phot = (
+            jnp.asarray(self.phot)
+            if self.phot is not None
+            else jnp.ones((n_lambda,))
+        )
+
+        # Limb-darkening coefficients
         u1 = jnp.atleast_1d(u1)
         u2 = jnp.atleast_1d(u2)
         u_ld = jnp.column_stack([u1, u2])
@@ -438,25 +616,28 @@ def transit_model(self, t0, period, rp, a, inclination,
         X = -r * jnp.cos(omega + true_anomaly)
         Y = -r * jnp.sin(omega + true_anomaly) * jnp.cos(inclination)
 
-        photosphere = (1 - f_S[..., 0].sum(axis=1)) * self.phot[None, :]
-
-        spot_coverages, spot_spectra = jnp.broadcast_arrays(
-            f_S[..., 0], self.spectrum[None, ...]
-        )
+        # Total fractional coverage by *any* spot
+        coverage_sum = f_S[..., 0].sum(axis=1)
+        # (n_times, 1)
 
-        time_series_spectrum = jnp.squeeze(
-            # photospheric component:
-            photosphere +
+        # Spot-weighted contrast
+        spot_weighted_contrast = jnp.sum(
+            f_S * contrast, axis=1
+        )[..., 0]
+        # (n_times, n_lambda)
 
-            # sum of the active region components:
-            jnp.sum(spot_coverages * spot_spectra, axis=1)
-        )
+        flux_factor = (1.0 - coverage_sum) + spot_weighted_contrast
+        time_series_spectrum = phot[None, :] * flux_factor
+        # (n_times, n_lambda)
 
-        rp = jnp.broadcast_to(rp, self.wavelength.shape)
+        rp = jnp.broadcast_to(rp, (n_lambda,))
+        # (n_lambda,)
 
-        # if one pair of limb-darkening coefficients are given,
-        # broadcast up to the shape of `rp`
+        # Broadcast limb-darkening if needed
         u_ld = u_ld * jnp.ones((rp.shape[0], 1))
+        # (n_lambda, 2)
+
+        # Uncontaminated transit model
         transit = vmap(
             lambda rp, u: jaxoplanet.core.light_curve(
                 u=u, b=jnp.hypot(X, Y), r=rp
@@ -464,7 +645,7 @@ def transit_model(self, t0, period, rp, a, inclination,
         )(rp, u_ld)
 
         contaminated_transit = (
-            time_series_spectrum - jnp.abs(transit) * self.phot[None, :]
+            time_series_spectrum - jnp.abs(transit) * phot[None, :]
         ) / time_series_spectrum
 
         t_ind = jnp.argmin(jnp.abs(self.times - t0))
@@ -559,6 +740,15 @@ def _area_union_per_time(
         beta = jnp.atleast_1d(beta)
         angle = jnp.atleast_1d(angle)
 
+        # If there are no active regions (no spots), return an empty
+        # (0, n_mc) array and skip the per-spot scan entirely. This
+        # prevents index errors when `occultation_possible` is empty.
+        if (
+            x0_ellipse.size == 0 or
+            getattr(occultation_possible, "size", 0) == 0
+        ):
+            return jnp.zeros((0, self.n_mc), dtype=bool)
+
         @jit
         def find_overlap(k):
             # find overlap between the planet and the elliptical region
@@ -623,7 +813,8 @@ def plot_star(self, t0, rp, a, inclination,
         if ax is None:
             ax = plt.gca()
 
-        log_temps = np.log10(self.temperature)
+        if self.temperature is not None:
+            log_temps = np.log10(self.temperature)
 
         def temp_cmap(x):
             return to_hex(
@@ -633,14 +824,54 @@ def temp_cmap(x):
                 )
             )
 
-        star = plt.Circle((0, 0), 1, color=to_hex(temp_cmap(self.T_eff)))
+        if self.T_eff is None:
+            color = 'white'
+        else:
+            color = to_hex(temp_cmap(self.T_eff))
+
+        star = plt.Circle((0, 0), 1, facecolor=color, edgecolor='k')
         ax.add_patch(star)
         ax.set(xlim=[-1.05, 1.05], ylim=[-1.05, 1.05])
 
-        squeezed_coords = list(map(
-            jnp.squeeze, self.spot_coords(times=jnp.array([t0]), t0_rot=t0_rot)
-        ))
-        for i, (x, y, z, _, _, _, _, angle) in enumerate(zip(*squeezed_coords)):
+        (
+            spot_position_x,
+            spot_position_y,
+            spot_position_z,
+            _,
+            _,
+            angle,
+            _,
+            contrast,
+        ) = self.spot_coords(times=jnp.array([t0, ]), t0_rot=t0_rot)
+
+        # Slice out that time and drop the singleton trailing dims.
+        # After this, each has shape (n_spots,).
+        spot_x = spot_position_x[0, :, 0, 0]
+        spot_y = spot_position_y[0, :, 0, 0]
+        spot_z = spot_position_z[0, :, 0, 0]
+        # contrast is usually (1, n_spots, n_lambda, 1)
+        # for plotting, pick λ = 0
+        contr = contrast[0, :, 0, 0]
+
+        # Convert to NumPy for plotting (avoids JAX tracing issues)
+        spot_x_np = np.asarray(spot_x)
+        spot_y_np = np.asarray(spot_y)
+        spot_z_np = np.asarray(spot_z)
+        contr_np = np.asarray(contr)
+
+        for i in range(self.n_spots):
+            x = spot_x_np[i]
+            y = spot_y_np[i]
+            z = spot_z_np[i]
+            c = contr_np[i]
+
+            if len(self.temperature) == 0:
+                color = 'k'
+                alpha = 1-c
+            else:
+                color = temp_cmap(self.temperature[i])
+                alpha = 1
+
             if z < 0:
                 rsq = x ** 2 + y ** 2
 
@@ -648,12 +879,17 @@ def temp_cmap(x):
                 angle = -np.degrees(np.arctan2(y, x))
                 ell = Ellipse(
                     (y, x), width=multiply_radii * 2 * self.rad[i],
-                    height=multiply_radii * 2 * self.rad[i] * short, angle=angle,
-                    facecolor=temp_cmap(self.temperature[i]), edgecolor='k'
+                    height=multiply_radii * 2 * self.rad[i] * short,
+                    angle=angle, facecolor=color, alpha=alpha, edgecolor='k'
                 )
                 ax.add_patch(ell)
 
-                if annotate:
+                if annotate and len(self.temperature) == 0:
+                    ax.annotate(
+                        f"{i+1}: {c:.2f}", (y, x),
+                        va='center', ha='center', fontsize=6
+                    )
+                elif annotate:
                     ax.annotate(
                         f"{i+1}: {int(self.temperature[i])} K", (y, x),
                         va='center', ha='center', fontsize=6

From 09f08ed0548b25afd2a53d4724a5bbe4a3c23e42 Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:23:04 -0500
Subject: [PATCH 06/10] Fixed orbital geometry calculations and enforced
 transit chord below stellar equator

---
 fleck/jax.py | 38 +++++++++++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index a035783..b26e9f3 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -603,8 +603,29 @@ def transit_model(self, t0, period, rp, a, inclination,
 
         f_S = rad ** 2 * mu * (spot_position_z < 0.).astype(int)
 
-        # compute the transit model
-        mean_anomaly = 2 * np.pi * (self.times - t0) / period
+        # ------------------------------------------------------------------
+        # Orbital geometry (independent of spots)
+        # ------------------------------------------------------------------
+        ecc = jnp.asarray(ecc)
+        omega = jnp.asarray(omega)
+
+        # Index of time sample closest to mid-transit
+        t_ind = jnp.argmin(jnp.abs(self.times - t0))
+
+        # True anomaly at (inferior) conjunction in this geometry
+        f_tr = 0.5 * jnp.pi - omega
+
+        # Convert f_tr -> E_tr -> M_tr:
+        #   tan(f/2) = sqrt((1+e)/(1-e)) * tan(E/2)
+        #   M = E - e sin(E)
+        sqrt_factor = jnp.sqrt((1.0 - ecc) / (1.0 + ecc + 1e-15))
+        E_tr = 2.0 * jnp.arctan(jnp.tan(0.5 * f_tr) * sqrt_factor)
+        M_tr = E_tr - ecc * jnp.sin(E_tr)
+
+        # Mean anomaly as a function of time with t0 at mid-transit
+        mean_anomaly = (
+            2.0 * jnp.pi * (self.times - t0) / period + M_tr
+        )
         true_anomaly = jnp.arctan2(
             *jaxoplanet.core.kepler(M=mean_anomaly, ecc=ecc)
         )
@@ -616,6 +637,18 @@ def transit_model(self, t0, period, rp, a, inclination,
         X = -r * jnp.cos(omega + true_anomaly)
         Y = -r * jnp.sin(omega + true_anomaly) * jnp.cos(inclination)
 
+        # --------------------------------------------------------------
+        # Enforce that the transit chord lies below the stellar equator
+        # in the (X, Y) plane used for spot occultations:
+        #
+        # Choose the branch with Y(t0) <= 0 by reflecting across the
+        # equator if necessary. This does NOT change b(t) or the
+        # light curve because b = hypot(X, Y) is sign-invariant.
+        # --------------------------------------------------------------
+        Y_mid = Y[t_ind]
+        sign_Y = jnp.where(Y_mid > 0., -1., 1.)
+        Y = sign_Y * Y
+
         # Total fractional coverage by *any* spot
         coverage_sum = f_S[..., 0].sum(axis=1)
         # (n_times, 1)
@@ -648,7 +681,6 @@ def transit_model(self, t0, period, rp, a, inclination,
             time_series_spectrum - jnp.abs(transit) * phot[None, :]
         ) / time_series_spectrum
 
-        t_ind = jnp.argmin(jnp.abs(self.times - t0))
         uncontaminated_max_depth = - transit[t_ind]
         contaminated_max_depth = (
             contaminated_transit.max(0) - contaminated_transit[t_ind]

From 30105228d30d00a2573a09069197c24f819886ea Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:23:13 -0500
Subject: [PATCH 07/10] Add omega parameter to plot_star method for argument of
 periapse

---
 fleck/jax.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index b26e9f3..b478ae9 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -813,7 +813,7 @@ def spot_step(carry, k, occultation_possible=occultation_possible):
         return monte_carlo_occulted_area
 
     def plot_star(self, t0, rp, a, inclination,
-                  ecc=0, t0_rot=0, multiply_radii=1,
+                  ecc=0, omega=jnp.pi/2, t0_rot=0, multiply_radii=1,
                   ax=None, annotate=False):
         """
         Plot a 2D representation of the star and transit chord.
@@ -830,6 +830,8 @@ def plot_star(self, t0, rp, a, inclination,
             Planetary orbital inclination [radians]
         ecc : float
             Orbital eccentricity, default is zero.
+        omega : float
+            Argument of periapse [radians], default is :math:`\\pi/2`.
         t0_rot : float
             Zero-point in time for stellar rotation, default is zero
         multiply_radii : float
@@ -930,7 +932,7 @@ def temp_cmap(x):
         ax.set_aspect('equal')
 
         b = (a * np.cos(inclination) * (1 - ecc ** 2) /
-             (1 + ecc * np.sin(np.pi / 2)))
+             (1 + ecc * np.sin(omega)))
 
         if hasattr(rp, '__len__'):
             rp = rp.mean()

From 270e0519dd523247a279d0b5b37d0b377d516a79 Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Sat, 29 Nov 2025 11:25:51 -0500
Subject: [PATCH 08/10] Fix color mapping in ActiveStar class by using T_eff
 for normalization

Resolves #23
---
 fleck/jax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index b478ae9..6688768 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -854,7 +854,7 @@ def temp_cmap(x):
             return to_hex(
                 plt.cm.YlOrRd_r(
                     (np.log10(x) - min(log_temps)) /
-                    (max(log_temps) - min(log_temps)) * 0.6 + 0.4
+                    (self.T_eff - min(log_temps)) * 0.6 + 0.4
                 )
             )
 

From 7b587e9d3c887d2a6eec5491770831abf1470246 Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Thu, 8 Oct 2026 15:47:54 -0400
Subject: [PATCH 09/10] Fix JAX active-region inputs and temperature-based flux
 normalization

Accept scalar and array contrasts and normalize spectra across spot
and wavelength axes. Initialize the first added spot correctly and
keep geometry, spectra, contrasts, and temperature labels aligned
when mixing spot definitions.

Derive a blackbody photosphere when none is supplied, while preserving
explicit photospheric spectra. Handle missing and equal temperatures
in plotting and normalize temperature colors consistently in log space.

Update examples and STSP tests to the signed-latitude convention,
removing the rotation test's compensating half-period offset.

Add regression coverage for input broadcasting, mixed spot definitions,
blackbody contrast ratios, plotting, and JAX differentiation.
---
 docs/fleck/jax.rst      |  18 +++--
 fleck/jax.py            | 156 ++++++++++++++++++++--------------------
 fleck/tests/test_jax.py | 136 ++++++++++++++++++++++++++++++++++-
 3 files changed, 224 insertions(+), 86 deletions(-)

diff --git a/docs/fleck/jax.rst b/docs/fleck/jax.rst
index 85e995c..76b9e3c 100644
--- a/docs/fleck/jax.rst
+++ b/docs/fleck/jax.rst
@@ -201,12 +201,22 @@ that we will observe, the stellar inclination, and spectrum of the stellar photo
 
 We can add active regions to the star with `~fleck.jax.ActiveStar.add_spot`
 
+Spot latitudes are signed angles in radians, from ``-pi/2`` to ``pi/2``;
+zero is the stellar equator. To convert a colatitude, use
+``latitude = pi/2 - colatitude``.
+
+Instead of a spectrum, you can provide a scalar or wavelength-dependent
+``contrast`` to ``add_spot``. A contrast below one describes a dark spot;
+a contrast above one describes a bright region. Supplying only ``temperature``
+uses a blackbody spot spectrum. If no photospheric spectrum is provided,
+``T_eff`` and the wavelength grid define a blackbody photosphere.
+
 .. code-block:: python
 
     # add a cool spot:
     active_star.add_spot(
         lon=-0.2,  # [rad]
-        lat=1.65,  # [rad]
+        lat=np.pi/2 - 1.65,  # signed latitude [rad]
         rad=0.15,  # [R_star]
         spectrum=cool.flux.value,
         temperature=cool.meta['PHXTEFF']
@@ -215,7 +225,7 @@ We can add active regions to the star with `~fleck.jax.ActiveStar.add_spot`
     # add a hot spot:
     active_star.add_spot(
         lon=0.95,
-        lat=1.75,
+        lat=np.pi/2 - 1.75,
         rad=0.08,
         spectrum=hot.flux.value,
         temperature=hot.meta['PHXTEFF']
@@ -376,7 +386,7 @@ code cell above, makes a lot of plotting calls to visualize the results of
     # add a cool spot:
     active_star.add_spot(
         lon=-0.2,  # [rad]
-        lat=1.65,  # [rad]
+        lat=np.pi/2 - 1.65,  # signed latitude [rad]
         rad=0.15,  # [R_star]
         spectrum=cool.flux.value,
         temperature=cool.meta['PHXTEFF']
@@ -385,7 +395,7 @@ code cell above, makes a lot of plotting calls to visualize the results of
     # add a hot spot:
     active_star.add_spot(
         lon=0.95,
-        lat=1.75,
+        lat=np.pi/2 - 1.75,
         rad=0.08,
         spectrum=hot.flux.value,
         temperature=hot.meta['PHXTEFF']
diff --git a/fleck/jax.py b/fleck/jax.py
index 6688768..0eb67d3 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -71,7 +71,9 @@ def __init__(
         wavelength : array
             Wavelength for each flux observation in ``phot`` [meters]
         phot : array
-            Photospheric flux at each ``wavelength``.
+            Photospheric flux at each ``wavelength``. If omitted, use a
+            blackbody when ``T_eff`` and ``wavelength`` are provided,
+            otherwise use unity.
         P_rot : float
             Stellar rotation period
         contrast : float or array
@@ -84,9 +86,10 @@ def __init__(
         self.lon = jnp.array(lon)
         self.lat = jnp.array(lat)
         self.rad = jnp.array(rad)
-        self.spectrum = jnp.array(spectrum)
+        self.spectrum = jnp.array([]) if spectrum is None else jnp.array(spectrum)
         self.T_eff = T_eff
-        self.temperature = jnp.array(temperature)
+        self.temperature = (jnp.array([]) if temperature is None
+                            else jnp.atleast_1d(jnp.asarray(temperature)))
         self.inclination = inclination
         # Allow wavelength/phot to remain optional, but store them as
         # JAX arrays when provided.
@@ -96,10 +99,7 @@ def __init__(
 
         # Raw contrast as passed by the user (may be scalar, 1D, 2D, or
         # empty). We normalise shapes after dimension inference.
-        if contrast is None or getattr(contrast, "size", 0) == 0:
-            self.contrast = jnp.array([])
-        else:
-            self.contrast = jnp.array(contrast)
+        self.contrast = jnp.array([]) if contrast is None else jnp.array(contrast)
 
         # ------------------------------------------------------------------
         # Dimension bookkeeping: infer n_times, n_spots, n_lambda
@@ -152,6 +152,11 @@ def __init__(
             # Absolute last resort: a single "white-light" wavelength bin
             self.n_lambda = 1
 
+        has_temperatures = self.temperature.size != 0
+        self.temperature = (jnp.broadcast_to(self.temperature, (self.n_spots,))
+                            if has_temperatures
+                            else jnp.full((self.n_spots,), jnp.nan))
+
         # Normalise the supplied contrast argument into something at least
         # 2D where possible.
         if self.contrast.size != 0:
@@ -169,22 +174,32 @@ def __init__(
             # If contrast is already 2D+ we assume the user passed
             # (n_spots, n_lambda) and leave it as-is.
 
-        # If we have no phot, set it to unity
-        self.phot = (jnp.asarray(self.phot) if self.phot is not None
-                     else jnp.ones((self.n_lambda,)))
-
-        # If an explicit contrast was not provided but we do have a
-        # spectrum and photosphere, derive the contrast now.
-        if self.contrast.size == 0 and self.spectrum.size != 0:
-            derived_contrast = self.spectrum / self.phot[None, :]
-            if derived_contrast.ndim == 1:
-                derived_contrast = derived_contrast[None, :]
-            self.contrast = derived_contrast
-
-        # If we have a contrast and a photosphere but no spectrum,
-        # construct a spectrum that is consistent with both.
-        if self.spectrum.size == 0 and self.contrast.size != 0:
+        if self.phot is None:
+            if self.T_eff is not None and self.wavelength is not None:
+                self.phot = self._blackbody(self.wavelength, self.T_eff)
+            else:
+                self.phot = jnp.ones((self.n_lambda,))
+        self.phot = jnp.broadcast_to(self.phot, (self.n_lambda,))
+
+        # Keep spectral arrays aligned along spot and wavelength axes.
+        shape = (self.n_spots, self.n_lambda)
+        if self.contrast.size != 0:
+            self.contrast = jnp.broadcast_to(self.contrast, shape)
             self.spectrum = self.contrast * self.phot[None, :]
+        else:
+            if self.spectrum.size == 0 and has_temperatures:
+                if self.wavelength is None:
+                    raise ValueError("Temperature-based spots require wavelengths.")
+                self.spectrum = self._blackbody(
+                    self.wavelength[None, :], self.temperature[:, None])
+            if self.spectrum.size != 0:
+                if self.spectrum.ndim == 1 and self.n_lambda == 1:
+                    self.spectrum = self.spectrum[:, None]
+                self.spectrum = jnp.broadcast_to(self.spectrum, shape)
+                self.contrast = self.spectrum / self.phot[None, :]
+            else:
+                self.contrast = jnp.zeros(shape)
+                self.spectrum = self.contrast * self.phot[None, :]
 
     def tree_flatten(self):
         children = (
@@ -420,58 +435,38 @@ def add_spot(self, lon, lat, rad, contrast=None, temperature=None,
         """
         if contrast is not None:
             # User-supplied contrast: convert to spectrum
-            contrast = jnp.array(contrast)
-            if contrast.ndim == 0:
-                # Scalar contrast: broadcast to all wavelengths
-                contrast = contrast[None] * jnp.ones(self.n_lambda)
-            elif contrast.ndim == 1:
-                # 1D contrast are defined on the wavelength grid
-                contrast = contrast[None, :]
+            contrast = jnp.broadcast_to(jnp.asarray(contrast).squeeze(),
+                                       (self.n_lambda,))
             # Derive the spectrum from the contrast
-            spectrum = contrast * self.phot[None, :]
+            spectrum = contrast * self.phot
         elif spectrum is not None:
             # User-supplied spectrum: ensure correct shape
-            spectrum = jnp.array(spectrum)
-            if spectrum.ndim == 0:
-                # Scalar spectrum: broadcast to all wavelengths
-                spectrum = spectrum[None] * jnp.ones(self.n_lambda)
-            elif spectrum.ndim == 1:
-                # 1D spectra are defined on the wavelength grid
-                spectrum = spectrum[None, :]
+            spectrum = jnp.broadcast_to(jnp.asarray(spectrum).squeeze(),
+                                       (self.n_lambda,))
             # Derive the contrast from the spectrum
-            contrast = spectrum / self.phot[None, :]
-        elif spectrum is None and temperature is not None:
+            contrast = spectrum / self.phot
+        elif temperature is not None:
             # No contrast or spectrum, but a temperature: compute a blackbody
+            if self.wavelength is None:
+                raise ValueError("Temperature-based spots require wavelengths.")
             spectrum = self._blackbody(self.wavelength, temperature)
-            contrast = spectrum / self.phot[None, :]
-            if contrast.ndim == 1:
-                contrast = contrast[None, :]
-
-        for attr, new_value in zip(
-            "lon, lat, rad, spectrum, temperature, contrast".split(', '),
-            [lon, lat, rad, spectrum, temperature, contrast],
-        ):
-
-            prop = getattr(self, attr)
-
-            if not hasattr(new_value, 'ndim'):
-                new_value = jnp.array([new_value])
-
-            if prop is not None and getattr(prop, "size", 0) != 0:
-                if prop.ndim > 1 or (
-                    len(prop) > 1 and len(prop) == len(new_value)
-                ):
-                    new_value = jnp.vstack([prop, new_value])
-                else:
-                    new_value = jnp.concatenate([prop, new_value])
-
-                setattr(self, attr, new_value)
-
-        # Update n_spots
-        self.n_spots += 1
+            contrast = spectrum / self.phot
+        else:
+            contrast = jnp.zeros((self.n_lambda,))
+            spectrum = contrast * self.phot
+
+        for attr, value in (("lon", lon), ("lat", lat), ("rad", rad),
+                            ("temperature", jnp.nan if temperature is None
+                             else temperature)):
+            setattr(self, attr, jnp.concatenate(
+                [jnp.ravel(getattr(self, attr)), jnp.asarray(value).reshape(1)]))
+        self.spectrum = jnp.concatenate([self.spectrum, spectrum[None, :]], axis=0)
+        self.contrast = jnp.concatenate([self.contrast, contrast[None, :]], axis=0)
+        self.n_spots = int(self.rad.size)
 
+    @staticmethod
     @jit
-    def _blackbody(self, wavelength_meters, temperature):
+    def _blackbody(wavelength_meters, temperature):
         """
         Compute a blackbody spectrum.
         """
@@ -847,16 +842,19 @@ def plot_star(self, t0, rp, a, inclination,
         if ax is None:
             ax = plt.gca()
 
-        if self.temperature is not None:
-            log_temps = np.log10(self.temperature)
+        temperatures = np.asarray(self.temperature)
+        known_temperatures = np.isfinite(temperatures) & (temperatures > 0)
+        color_temperatures = temperatures[known_temperatures]
+        if self.T_eff is not None:
+            color_temperatures = np.append(color_temperatures, self.T_eff)
+        log_temps = np.log10(color_temperatures)
+        log_min = log_temps.min() if log_temps.size else 0.
+        log_max = log_temps.max() if log_temps.size else 0.
 
         def temp_cmap(x):
-            return to_hex(
-                plt.cm.YlOrRd_r(
-                    (np.log10(x) - min(log_temps)) /
-                    (self.T_eff - min(log_temps)) * 0.6 + 0.4
-                )
-            )
+            fraction = ((np.log10(x) - log_min) / (log_max - log_min)
+                        if log_max > log_min else 1.)
+            return to_hex(plt.cm.YlOrRd_r(fraction * 0.6 + 0.4))
 
         if self.T_eff is None:
             color = 'white'
@@ -899,11 +897,11 @@ def temp_cmap(x):
             z = spot_z_np[i]
             c = contr_np[i]
 
-            if len(self.temperature) == 0:
-                color = 'k'
-                alpha = 1-c
+            if not known_temperatures[i]:
+                color = 'k' if c <= 1 else 'white'
+                alpha = np.clip(abs(1-c), 0, 1)
             else:
-                color = temp_cmap(self.temperature[i])
+                color = temp_cmap(temperatures[i])
                 alpha = 1
 
             if z < 0:
@@ -918,7 +916,7 @@ def temp_cmap(x):
                 )
                 ax.add_patch(ell)
 
-                if annotate and len(self.temperature) == 0:
+                if annotate and not known_temperatures[i]:
                     ax.annotate(
                         f"{i+1}: {c:.2f}", (y, x),
                         va='center', ha='center', fontsize=6
diff --git a/fleck/tests/test_jax.py b/fleck/tests/test_jax.py
index ded831c..d6629ea 100644
--- a/fleck/tests/test_jax.py
+++ b/fleck/tests/test_jax.py
@@ -37,12 +37,12 @@ def test_rotation():
     )
 
     active_star.lon = jnp.radians(jnp.array([lon1, lon2]))
-    active_star.lat = jnp.radians(jnp.array([lat1, lat2])) - np.pi / 2
+    active_star.lat = jnp.radians(jnp.array([lat1, lat2]))
     active_star.rad = jnp.array([rad1, rad2])
     active_star.spectrum = jnp.array([0.7, 0.7])
 
     fleck_lc = active_star.rotation_model(
-        f0=1, t0_rot=0.5, u1=u_ld[0], u2=u_ld[1]
+        f0=1, t0_rot=0, u1=u_ld[0], u2=u_ld[1]
     )[0][:, 0, 0]
 
     # Assert matches STSP results to within 50 ppm:
@@ -80,7 +80,7 @@ def test_jax_transit():
     active_star = ActiveStar(
         times=times,
         lon=np.ravel(spot_lons.to_value(u.rad)),
-        lat=np.ravel(np.pi/2 + spot_lats.to_value(u.rad)),
+        lat=np.ravel(spot_lats.to_value(u.rad)),
         rad=np.ravel(spot_radii),
         inclination=inc_stellar.to_value(u.rad),
         phot=jnp.ones(10),
@@ -99,3 +99,133 @@ def test_jax_transit():
 
     assert np.max(np.abs(stsp_lc - jax_lc[:, 0]) / jax_lc[:, 0]) < 5.0e-4
     assert np.std(stsp_lc / jax_lc[:, 0] - 1) < 1.5e-4
+
+
+@pytest.fixture
+def star_inputs():
+    return dict(times=jnp.linspace(-0.05, 0.05, 11), lon=[0.], lat=[0.],
+                rad=[0.1], inclination=np.pi/2, wavelength=[1e-6],
+                phot=[1.], P_rot=10.)
+
+
+@pytest.mark.parametrize('contrast', [0.5, [0.5], np.array([0.5]),
+                                      jnp.array(0.5), jnp.array([0.5])])
+def test_contrast_inputs_produce_half_dark_spot(star_inputs, contrast):
+    star = ActiveStar(**star_inputs, contrast=contrast)
+    np.testing.assert_allclose(star.contrast, [[0.5]])
+    flux = np.asarray(star.rotation_model(f0=1)[0])
+    # At disk center a spot removes its area times its missing flux fraction.
+    assert flux[len(flux)//2, 0, 0] == pytest.approx(1 - 0.1**2*0.5,
+                                                   abs=1e-6)
+
+
+@pytest.mark.parametrize('contrast, expected', [
+    (0.5, [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]]),
+    ([0.4, 0.7], [[0.4, 0.4, 0.4], [0.7, 0.7, 0.7]]),
+    ([0.4, 0.5, 0.7], [[0.4, 0.5, 0.7], [0.4, 0.5, 0.7]]),
+])
+def test_contrast_broadcasts_over_spots_and_wavelengths(contrast, expected):
+    star = ActiveStar(times=[0.], lon=[0., 0.2], lat=[0., 0.1],
+                      rad=[0.1, 0.05], inclination=np.pi/2,
+                      wavelength=[1e-6, 2e-6, 3e-6], phot=[2., 3., 4.],
+                      contrast=contrast)
+    np.testing.assert_allclose(star.contrast, expected)
+    np.testing.assert_allclose(star.spectrum, np.asarray(expected)*[2., 3., 4.])
+    assert star.rotation_model(f0=1)[0].shape == (1, 3, 1)
+
+
+def test_add_first_spot_and_mix_temperature_and_contrast():
+    star = ActiveStar(times=jnp.linspace(-0.05, 0.05, 11),
+                      inclination=np.pi/2, wavelength=[1e-6, 2e-6],
+                      T_eff=5000.)
+    original_phot = np.asarray(star.phot).copy()
+    star.add_spot(lon=0., lat=0., rad=0.1, contrast=0.5)
+    star.add_spot(lon=0.2, lat=-0.1, rad=0.05, temperature=4000.)
+    star.add_spot(lon=-0.2, lat=0.1, rad=0.05, spectrum=0.8*original_phot)
+
+    assert star.n_spots == 3
+    np.testing.assert_allclose(star.lon, [0., 0.2, -0.2])
+    np.testing.assert_allclose(star.lat, [0., -0.1, 0.1])
+    np.testing.assert_allclose(star.rad, [0.1, 0.05, 0.05])
+    np.testing.assert_allclose(star.contrast[[0, 2], :], [[0.5, 0.5], [0.8, 0.8]])
+    np.testing.assert_allclose(star.phot, original_phot)
+    assert star.spectrum.shape == (3, 2)
+    assert star.contrast.shape == (3, 2)
+    assert np.isnan(np.asarray(star.temperature)[[0, 2]]).all()
+    assert star.temperature[1] == 4000.
+    assert np.isfinite(star.rotation_model(f0=1)[0]).all()
+
+
+@pytest.mark.parametrize('construct_spot', [True, False])
+def test_temperature_spot_uses_blackbody_photosphere(construct_spot):
+    wavelength = np.array([1e-6, 2e-6])
+    kwargs = dict(times=[0.], inclination=np.pi/2, wavelength=wavelength,
+                  T_eff=5000.)
+    if construct_spot:
+        star = ActiveStar(**kwargs, lon=[0.], lat=[0.], rad=[0.1],
+                          temperature=[4000.])
+    else:
+        star = ActiveStar(**kwargs)
+        star.add_spot(lon=0., lat=0., rad=0.1, temperature=4000.)
+    hc_over_k = 6.62607015e-34*299792458./1.380649e-23
+    expected = (np.expm1(hc_over_k/(wavelength*5000.)) /
+                np.expm1(hc_over_k/(wavelength*4000.)))
+    np.testing.assert_allclose(star.contrast[0], expected, rtol=1e-5)
+    assert np.all((star.contrast > 0) & (star.contrast < 1))
+
+
+def test_explicit_photosphere_is_preserved(star_inputs):
+    star = ActiveStar(**star_inputs, T_eff=5000., contrast=0.5)
+    np.testing.assert_array_equal(star.phot, [1.])
+    np.testing.assert_allclose(star.spectrum, [[0.5]])
+
+
+@pytest.mark.parametrize('temperature', [None, 5000.])
+@pytest.mark.parametrize('contrast', [0.5, 1.5])
+def test_plot_contrast_only_spot_without_temperature(star_inputs, temperature,
+                                                    contrast):
+    import matplotlib.pyplot as plt
+
+    star = ActiveStar(**star_inputs, contrast=contrast, T_eff=temperature)
+    fig, ax = plt.subplots()
+    try:
+        star.plot_star(0., 0.1, 10., np.deg2rad(89), ax=ax, annotate=True)
+        fig.canvas.draw()
+        assert len(ax.patches) == 2
+        assert ax.patches[1].get_alpha() == pytest.approx(0.5)
+        assert ax.texts[0].get_text() == f'1: {contrast:.2f}'
+    finally:
+        plt.close(fig)
+
+
+@pytest.mark.parametrize('temperatures', [[4000., 5000.], [5000., 5000.],
+                                         [4000., 6000.]])
+def test_plot_temperature_colors_follow_temperature(temperatures):
+    import matplotlib.pyplot as plt
+
+    star = ActiveStar(times=[0.], lon=[-0.2, 0.2], lat=[0., 0.],
+                      rad=[0.1, 0.1], inclination=np.pi/2,
+                      wavelength=[1e-6], phot=[1.], contrast=0.5,
+                      T_eff=5000., temperature=temperatures)
+    fig, ax = plt.subplots()
+    try:
+        star.plot_star(0., 0.1, 10., np.deg2rad(89), ax=ax, annotate=True)
+        fig.canvas.draw()
+        colors = np.array([patch.get_facecolor() for patch in ax.patches])
+        assert np.isfinite(colors).all()
+        if temperatures[0] == temperatures[1]:
+            np.testing.assert_array_equal(colors[0], colors[1])
+        else:
+            assert not np.array_equal(colors[0], colors[1])
+            assert not np.array_equal(colors[1], colors[2])
+        assert len(ax.texts) == 2
+    finally:
+        plt.close(fig)
+
+
+def test_contrast_is_differentiable(star_inputs):
+    def center_flux(contrast):
+        star = ActiveStar(**star_inputs, contrast=contrast)
+        return star.rotation_model(f0=1)[0][5, 0, 0]
+
+    assert float(jax.grad(center_flux)(0.5)) == pytest.approx(0.1**2, abs=1e-6)

From ba4695a01b7f3d266ab3f97c967f0c4b49f05e5e Mon Sep 17 00:00:00 2001
From: Taylor J Bell <tbell@stsci.edu>
Date: Thu, 8 Oct 2026 15:51:27 -0400
Subject: [PATCH 10/10] Fix Flake8 operator line breaks in JAX calculations

---
 fleck/jax.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/fleck/jax.py b/fleck/jax.py
index 0eb67d3..6c9dc63 100644
--- a/fleck/jax.py
+++ b/fleck/jax.py
@@ -731,8 +731,8 @@ def time_step(
             (1. - contrast) *
             jnp.expand_dims(frac_occulted_per_time_per_spot, axis=(2, 3))
         )
-        scaled_occultation = ((1. - contaminated_transit)
-                              * jnp.sum(occultation, axis=1)[..., 0])
+        scaled_occultation = ((1. - contaminated_transit) *
+                              jnp.sum(occultation, axis=1)[..., 0])
 
         spectrum_at_transit = time_series_spectrum[t_ind]
 
@@ -1065,6 +1065,6 @@ def spectral_binning(y, all_x, all_y):
     min_ind = np.argwhere(all_y == y[0])[0, 0]
     max_ind = np.argwhere(all_y == y[-1])[0, 0]
     if max_ind > min_ind and y.shape == all_x[min_ind:max_ind + 1].shape:
-        return (np.trapz(y, all_x[min_ind:max_ind + 1])
-                / (all_x[max_ind] - all_x[min_ind]))
+        return (np.trapz(y, all_x[min_ind:max_ind + 1]) /
+                (all_x[max_ind] - all_x[min_ind]))
     return np.nan