fix: check spectrum input and shape of outputs

fix: check spectrum input and shape of outputs

Author: lgrcia

.github/workflows/ci.yml

@@ -1,15 +1,15 @@
 name: spotter/ci
 
 on:
-  merge_group:
-  workflow_dispatch:
-  pull_request:
-  release:
-    types: [published]
+  push:
     branches:
       - main
     tags:
       - "*"
+  pull_request:
+  merge_group:
+  release:
+    types: [published]
 
 jobs:
   tests:
@@ -70,7 +70,7 @@ jobs:
   publish:
     environment:
       name: pypi
-      url: https://pypi.org/p/spotter
+      url: https://pypi.org/p/fitsm
     permissions:
       id-token: write
     needs: [tests, build]
@@ -79,6 +79,6 @@ jobs:
     steps:
       - uses: actions/download-artifact@v4
         with:
-          name: artifact
+          name: dist
           path: dist
-      - uses: pypa/gh-action-pypi-publish@v1.8.11
+      - uses: pypa/gh-action-pypi-publish@v1.12.4

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "spotter"
-version = "0.0.8"
+version = "0.0.9"
 description = "Approximate forward models of fluxes and spectra time-series of non-uniform stars."
 authors = [{ name = "Lionel Garcia" }, { name = "Benjamin Rackham" }]
 license = "MIT"

spotter/core.py

@@ -533,7 +533,9 @@ def shifted_spectra(spectra, shift):
     return jnp.real(shifted)
 
 
-def integrated_spectrum(N, theta, phi, period, radius, wv, spectra, phase, inc, normalize = True):
+def integrated_spectrum(
+    N, theta, phi, period, radius, wv, spectra, phase, inc, normalize=True
+):
     """
     Compute the integrated spectrum of a rotating star.
 
@@ -565,17 +567,23 @@ def integrated_spectrum(N, theta, phi, period, radius, wv, spectra, phase, inc,
     """
     spectra = jnp.atleast_2d(spectra)
     if period is None:
-        spectra_shifted = spectra
+        mask, projected, limb = mask_projected_limb(vec(N), 0.0, inc=inc)
+        limb_geometry = projected * mask * limb
+        spectra_shifted = spectra.T
+
     else:
         mask, projected, limb = mask_projected_limb(vec(N), phase, inc=inc)
         w_shift = doppler_shift(theta, phi, period, radius, phase)
         dw = wv[1] - wv[0]
         shift = w_shift[:, None] * wv / dw
         limb_geometry = projected * mask * limb
         spectra_shifted = shifted_spectra(spectra.T, shift)
-        
-        if normalize:
-            integrated_spec =  jnp.sum(spectra_shifted * limb_geometry[:, None], 0) / jnp.sum(limb_geometry[:, None] * spectra.T)
-        else:
-            integrated_spec =  jnp.sum(spectra_shifted * limb_geometry[:, None], 0)
-        return integrated_spec
+
+    if normalize:
+        integrated_spec = jnp.sum(
+            spectra_shifted * limb_geometry[:, None], 0
+        ) / jnp.sum(limb_geometry[:, None] * spectra.T)
+    else:
+        integrated_spec = jnp.sum(spectra_shifted * limb_geometry[:, None], 0)
+
+    return integrated_spec

spotter/doppler.py

@@ -33,8 +33,18 @@ def spectrum(star: Star, time: float, normalize: bool = True) -> ArrayLike:
     spectrum : ndarray
         Integrated spectrum.
     """
+    if star.wv is None:
+        raise ValueError("Star.wv must be set.")
+
+    if star.wv.shape[0] != star.y.shape[0]:
+        raise ValueError(
+            "The star spectrum (Star.y) must have the same number of wavelength as the "
+            f"wavelengths provided (Star.wv).\nFound Star.wv.shape[0] = {star.wv.shape[0]} "
+            f"and Star.y.shape[0] = {star.y.shape[0]}"
+        )
+
     phi, theta = hp.pix2ang(star.sides, range(hp.nside2npix(star.sides)))
-    
+
     def impl(star, time):
         return core.integrated_spectrum(
             star.sides,
@@ -52,7 +62,6 @@ def impl(star, time):
     return jnp.vectorize(impl, excluded=(0,), signature="()->(n)")(star, time)
 
 
-
 def rv_design_matrix(star: Star, time: float) -> ArrayLike:
     """
     Compute the radial velocity design matrix for a Star.

tests/test_spectrum.py

@@ -0,0 +1,31 @@
+from spotter import Star, doppler
+import numpy as np
+import pytest
+
+
+def test_star_inputs():
+    wv = np.linspace(0, 1, 10)
+    star = Star.from_sides(10, wv=wv, period=1.0)
+    with pytest.raises(ValueError):
+        doppler.spectrum(star, [0])
+
+    star = Star.from_sides(10)
+    with pytest.raises(ValueError):
+        doppler.spectrum(star, [0])
+
+
+def test_spectrum():
+    wv = np.linspace(0, 1, 10)
+    star = Star.from_sides(10, wv=wv, period=1.0)
+    spectra = np.random.rand(len(wv), star.size)
+    star = star.set(y=spectra)
+    assert doppler.spectrum(star, [0]).shape == (1, len(wv))
+    assert doppler.spectrum(star, [0, 1, 2]).shape == (3, len(wv))
+
+
+def test_no_period():
+    wv = np.linspace(0, 1, 10)
+    star = Star.from_sides(10, wv=wv)
+    spectra = np.random.rand(len(wv), star.size)
+    star = star.set(y=spectra)
+    assert doppler.spectrum(star, [0]).shape == (1, len(wv))