feat: add n_cells method

python/mocpy/fmoc/fmoc.py

@@ -1,9 +1,8 @@
 import numpy as np
 from astropy import units as u
 
-from ..abstract_moc import AbstractMOC
-
 from .. import mocpy
+from ..abstract_moc import AbstractMOC
 
 __author__ = "Matthieu Baumann, Thomas Boch, Manon Marchand, François-Xavier Pineau"
 __copyright__ = "CDS, Centre de Données astronomiques de Strasbourg"
@@ -61,6 +60,33 @@ def max_order(self):
         depth = mocpy.get_fmoc_depth(self._store_index)
         return np.uint8(depth)
 
+    @classmethod
+    def n_cells(cls, depth):
+        """Get the number of cells for a given depth.
+
+        Parameters
+        ----------
+        depth : int
+            The depth. It is comprised between 0 and `~mocpy.fmoc.FrequencyMOC.MAX_ORDER`
+
+        Returns
+        -------
+        int
+            The number of cells at the given order
+
+        Examples
+        --------
+        >>> from mocpy import FrequencyMOC
+        >>> FrequencyMOC.n_cells(0)
+        2
+        """
+        if depth < 0 or depth > cls.MAX_ORDER:
+            raise ValueError(
+                f"The depth should be comprised between 0 and {cls.MAX_ORDER}, but {depth}"
+                " was provided.",
+            )
+        return mocpy.n_cells_fmoc(depth)
+
     def to_hz_ranges(self):
         """Return the Hertz ranges this `FrequencyMOC` contains, in Hertz.
 
@@ -514,8 +540,8 @@ def plot_frequencies(self, ax, color="blue", frequency_unit="Hz"):
                 " instead of 'frequency', see astropy.units for more information",
             )
 
-        from matplotlib.patches import Rectangle
         from matplotlib.collections import PatchCollection
+        from matplotlib.patches import Rectangle
 
         min_freq = self.min_freq.to(frequency_unit).value
         max_freq = self.max_freq.to(frequency_unit).value
@@ -568,8 +594,8 @@ def plot_wavelengths(self, ax, color="blue", length_unit="m"):
                 " instead of 'length', see astropy.units for more information",
             )
 
-        from matplotlib.patches import Rectangle
         from matplotlib.collections import PatchCollection
+        from matplotlib.patches import Rectangle
 
         # get default bonds
         min_lambda = self.max_freq.to(length_unit, equivalencies=u.spectral()).value

python/mocpy/moc/moc.py

@@ -1,6 +1,7 @@
+import contextlib
+import functools
 from io import BytesIO
 from urllib.parse import urlencode
-import functools
 
 import numpy as np
 from astropy import units as u
@@ -17,8 +18,6 @@
 from astropy.io import fits
 from astropy.utils.data import download_file
 
-import contextlib
-
 with contextlib.suppress(ImportError):
     from astropy_healpix import HEALPix
 
@@ -141,6 +140,33 @@ def max_order(self):
         depth = mocpy.get_smoc_depth(self._store_index)
         return np.uint8(depth)
 
+    @classmethod
+    def n_cells(cls, depth):
+        """Get the number of cells for a given depth.
+
+        Parameters
+        ----------
+        depth : int
+            The depth. It is comprised between 0 and `~mocpy.moc.MOC.MAX_ORDER`
+
+        Returns
+        -------
+        int
+            The number of cells at the given order
+
+        Examples
+        --------
+        >>> from mocpy import MOC
+        >>> MOC.n_cells(0)
+        12
+        """
+        if depth < 0 or depth > cls.MAX_ORDER:
+            raise ValueError(
+                f"The depth should be comprised between 0 and {cls.MAX_ORDER}, but {depth}"
+                " was provided.",
+            )
+        return mocpy.n_cells_smoc(depth)
+
     def split_count(self, include_indirect_neighbours=False):
         """
         Return the number of disjoint MOCs the given MOC contains.

python/mocpy/stmoc/stmoc.py

@@ -1,9 +1,8 @@
 import numpy as np
 
 from .. import MOC, mocpy
-from ..tmoc import TimeMOC, microseconds_to_times, times_to_microseconds
 from ..abstract_moc import AbstractMOC
-
+from ..tmoc import TimeMOC, microseconds_to_times, times_to_microseconds
 
 __author__ = "Matthieu Baumann, Thomas Boch, Manon Marchand, François-Xavier Pineau"
 __copyright__ = "CDS, Centre de Données astronomiques de Strasbourg"
@@ -57,6 +56,39 @@ def min_time(self):
         """Return STMOC min time."""
         return microseconds_to_times(mocpy.coverage_2d_min_time(self._store_index))
 
+    @classmethod
+    def n_cells(cls, depth, dimension):
+        """Get the number of cells for a given depth.
+
+        Parameters
+        ----------
+        depth : int
+            The depth. It is comprised between 0 and `~mocpy.moc.MOC.MAX_ORDER` if
+            dimension='space' and between 0 and `~mocpy.tmoc.TimeMOC.MAX_ORDER` if
+            dimension='time'.
+
+        dimension : str
+            Can be either 'time' or 'space'.
+
+        Returns
+        -------
+        int
+            The number of cells at the given order
+
+        Examples
+        --------
+        >>> from mocpy import STMOC
+        >>> STMOC.n_cells(0, dimension='space')
+        12
+        """
+        if dimension == "space":
+            return MOC.n_cells(depth)
+        if dimension == "time":
+            return TimeMOC.n_cells(depth)
+        raise ValueError(
+            f"Dimension should be either 'time' of 'space' but '{dimension}' was provided.",
+        )
+
     def is_empty(self):
         """Check whether the Space-Time coverage is empty."""
         return mocpy.is_empty(self._store_index)
@@ -194,8 +226,11 @@ def from_spatial_coverages(
         result : `~mocpy.stmoc.STMOC`
             The resulting Spatial-Time Coverage map.
         """
-        # times_start = times_start.jd.astype(np.float64)
-        # times_end = times_end.jd.astype(np.float64)
+        # accept also when there is a single spatial moc
+        times_start = np.atleast_1d(times_start)
+        times_end = np.atleast_1d(times_end)
+        spatial_coverages = np.atleast_1d(spatial_coverages)
+
         times_start = times_to_microseconds(times_start)
         times_end = times_to_microseconds(times_end)
 

python/mocpy/tests/test_fmoc.py

@@ -1,7 +1,8 @@
 import numpy as np
+import pytest
+from astropy import units as u
 
 from ..fmoc import FrequencyMOC
-from astropy import units as u
 
 
 def test_new_empty():
@@ -14,6 +15,17 @@ def test_max_order():
     assert fmoc.max_order == 32
 
 
+def test_n_cells():
+    assert FrequencyMOC.n_cells(0) == 2
+    with pytest.raises(
+        ValueError,
+        match=f"The depth should be comprised between 0 and {FrequencyMOC.MAX_ORDER}*",
+    ):
+        FrequencyMOC.n_cells(-1)
+        FrequencyMOC.n_cells(FrequencyMOC.MAX_ORDER + 1)
+    assert FrequencyMOC.n_cells(5) == 2 * FrequencyMOC.n_cells(4)
+
+
 def test_to_depth59_ranges():
     fmoc = FrequencyMOC.new_empty(FrequencyMOC.MAX_ORDER).complement()
     # 2^60 = 1152921504606846976

python/mocpy/tests/test_moc.py

@@ -1,15 +1,13 @@
-import pytest
 import copy
 import re
 
-import numpy as np
-
-from astropy.coordinates import SkyCoord, Angle
-from astropy.io.votable import parse_single_table
 import astropy.units as u
-from astropy.io import fits
-
 import cdshealpix
+import numpy as np
+import pytest
+from astropy.coordinates import Angle, SkyCoord
+from astropy.io import fits
+from astropy.io.votable import parse_single_table
 
 from ..moc import MOC, WCS
 
@@ -98,6 +96,17 @@ def test_to_depth29_ranges(isets):
     assert np.array_equal(l, r)
 
 
+def test_n_cells():
+    assert MOC.n_cells(0) == 12
+    with pytest.raises(
+        ValueError,
+        match=f"The depth should be comprised between 0 and {MOC.MAX_ORDER}*",
+    ):
+        MOC.n_cells(-2)
+        MOC.n_cells(MOC.MAX_ORDER + 1)
+    assert MOC.n_cells(6) == 4 * MOC.n_cells(5)
+
+
 def test_interval_set_union(isets):
     assert isets["a"].union(isets["b"]) == MOC.from_depth29_ranges(
         29,
@@ -416,13 +425,6 @@ def test_serialize_to_str(moc, expected):
         ("moc", False, "ascii", True),
     ],
 )
-def test_write(moc_from_json, filename, overwrite, format, os_error):  # noqa: A002
-    if os_error:
-        with pytest.raises(OSError):  # TODO add the match parameter of the exception
-            moc_from_json.save(filename, format=format, overwrite=overwrite)
-    else:
-        moc_from_json.save(filename, format=format, overwrite=overwrite)
-
 
 # --- TESTING MOC plot functions ---#
 def test_mpl_fill():
@@ -519,7 +521,7 @@ def test_degrade_to_order():
 
     max_depth = hst_moc.max_order
 
-    for order in reversed(range(0, max_depth)):
+    for order in reversed(range(max_depth)):
         hst_moc = hst_moc.degrade_to_order(order)
         assert hst_moc.sky_fraction <= 1.0
 
@@ -571,11 +573,8 @@ def mocs():
 
 
 def test_add_neighbours(mocs):
-    print(mocs["moc1"])
     mocs["moc1"].add_neighbours()
-    print(mocs["moc1"])
     assert mocs["moc1"] == mocs["moc1_increased"]
-
     mocs["moc2"].add_neighbours()
     assert mocs["moc2"] == mocs["moc2_increased"]
 
@@ -723,8 +722,8 @@ def test_from_valued_healpix_cells_bayestar():
     uniq = data["UNIQ"]
     probdensity = data["PROBDENSITY"]
 
-    import astropy_healpix as ah
     import astropy.units as u
+    import astropy_healpix as ah
 
     level, _ = ah.uniq_to_level_ipix(uniq)
     area = ah.nside_to_pixel_area(ah.level_to_nside(level)).to_value(u.steradian)
@@ -750,28 +749,3 @@ def test_from_valued_healpix_cells_bayestar_and_split():
     assert len(mocs) == 2
     for moc in mocs:
         assert moc.max_order == 11
-
-
-# --- TESTING new features ---#
-def test_moc_save_load_deser():
-    smoc = MOC.from_string("3/3 10 4/16-18 22 5/19-20 17/222 28/123456789 29/", "ascii")
-    smoc.to_string("ascii")
-    smoc_json = smoc.to_string("json")
-    smoc_bis = MOC.from_string(smoc_json, "json")
-    assert smoc == smoc_bis
-
-    smoc_bis = MOC.load("resources/MOC2.0/smoc.ascii.txt", "ascii")
-    assert smoc == smoc_bis
-
-    smoc_bis = MOC.load("resources/MOC2.0/SMOC.fits", "fits")
-    assert smoc == smoc_bis
-
-    smoc.save("resources/MOC2.0/smoc.py.test.fits", format="fits", overwrite=True)
-    smoc.save("resources/MOC2.0/smoc.py.test.json", format="json", overwrite=True)
-    smoc.save("resources/MOC2.0/smoc.py.test.ascii", format="ascii", overwrite=True)
-    smoc_bis = MOC.load("resources/MOC2.0/smoc.py.test.fits", "fits")
-    assert smoc == smoc_bis
-    smoc_bis = MOC.load("resources/MOC2.0/smoc.py.test.json", "json")
-    assert smoc == smoc_bis
-    smoc_bis = MOC.load("resources/MOC2.0/smoc.py.test.ascii", "ascii")
-    assert smoc == smoc_bis

python/mocpy/tests/test_stmoc.py

@@ -1,13 +1,13 @@
-from ..stmoc import STMOC
-from ..tmoc import TimeMOC
-from ..moc import MOC
-from astropy.time import Time, TimeDelta
-from astropy.table import Table
 import astropy.units as u
-
-import pytest
 import cdshealpix
 import numpy as np
+import pytest
+from astropy.table import Table
+from astropy.time import Time, TimeDelta
+
+from ..moc import MOC
+from ..stmoc import STMOC
+from ..tmoc import TimeMOC
 
 
 @pytest.fixture()
@@ -52,6 +52,16 @@ def stmoc_xmm_dr8():
     )
 
 
+def test_n_cells():
+    assert STMOC.n_cells(0, "time") == 2
+    assert STMOC.n_cells(0, "space") == 12
+    with pytest.raises(
+        ValueError,
+        match="Dimension should be either 'time' of 'space' but 'nothing' was provided.",
+    ):
+        STMOC.n_cells(0, "nothing")
+
+
 def test_serialization():
     decals = STMOC.from_fits("resources/STMOC/STMoc-DECaLS-g.fits")
     decals_bis = STMOC.load("resources/STMOC/STMoc-DECaLS-g.fits", format="fits")
@@ -216,28 +226,3 @@ def test_stmoc_from_spatial_coverages():
 
     expected_stmoc = STMOC.from_string("t59/1 60/1 61/8 s28/0")
     assert stmoc == expected_stmoc
-
-
-# --- TESTING new features ---#
-def test_stmoc_save_load_deser():
-    stmoc = STMOC.from_string("t61/1 3 5 s3/1-3 4/ t61/50 52 s4/25", "ascii")
-    stmoc_json = stmoc.to_string("json")
-
-    stmoc_bis = STMOC.from_string(stmoc_json, "json")
-    assert stmoc == stmoc_bis
-
-    stmoc_bis = STMOC.load("resources/MOC2.0/stmoc.ascii.txt", "ascii")
-    assert stmoc == stmoc_bis
-
-    stmoc_bis = STMOC.load("resources/MOC2.0/STMOC.fits", "fits")
-    assert stmoc == stmoc_bis
-
-    stmoc.save("resources/MOC2.0/stmoc.py.test.fits", format="fits", overwrite=True)
-    stmoc.save("resources/MOC2.0/stmoc.py.test.json", format="json", overwrite=True)
-    stmoc.save("resources/MOC2.0/stmoc.py.test.ascii", format="ascii", overwrite=True)
-    stmoc_bis = STMOC.load("resources/MOC2.0/stmoc.py.test.fits", "fits")
-    assert stmoc == stmoc_bis
-    stmoc_bis = STMOC.load("resources/MOC2.0/stmoc.py.test.json", "json")
-    assert stmoc == stmoc_bis
-    stmoc_bis = STMOC.load("resources/MOC2.0/stmoc.py.test.ascii", "ascii")
-    assert stmoc == stmoc_bis