From cf51f730fdf6e268ec26f8abb8cc983e4006173c Mon Sep 17 00:00:00 2001 From: MARCHAND MANON Date: Wed, 15 Nov 2023 17:35:38 +0100 Subject: [PATCH] feat: add n_cells method --- python/mocpy/fmoc/fmoc.py | 34 ++++++++++++++-- python/mocpy/moc/moc.py | 32 ++++++++++++++-- python/mocpy/stmoc/stmoc.py | 43 +++++++++++++++++++-- python/mocpy/tests/test_fmoc.py | 14 ++++++- python/mocpy/tests/test_moc.py | 62 +++++++++--------------------- python/mocpy/tests/test_stmoc.py | 49 ++++++++---------------- python/mocpy/tests/test_tmoc.py | 66 +++++++------------------------- python/mocpy/tmoc/tmoc.py | 39 +++++++++++++++---- src/lib.rs | 32 +++++++++++++++- 9 files changed, 222 insertions(+), 149 deletions(-) diff --git a/python/mocpy/fmoc/fmoc.py b/python/mocpy/fmoc/fmoc.py index 768b5e8b..44791933 100644 --- a/python/mocpy/fmoc/fmoc.py +++ b/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 diff --git a/python/mocpy/moc/moc.py b/python/mocpy/moc/moc.py index c21c34d6..b7429a31 100644 --- a/python/mocpy/moc/moc.py +++ b/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. diff --git a/python/mocpy/stmoc/stmoc.py b/python/mocpy/stmoc/stmoc.py index 8d0d8d8b..f5c82d03 100644 --- a/python/mocpy/stmoc/stmoc.py +++ b/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) diff --git a/python/mocpy/tests/test_fmoc.py b/python/mocpy/tests/test_fmoc.py index de11ae56..680629e6 100644 --- a/python/mocpy/tests/test_fmoc.py +++ b/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 diff --git a/python/mocpy/tests/test_moc.py b/python/mocpy/tests/test_moc.py index ea0cc381..291ce7af 100644 --- a/python/mocpy/tests/test_moc.py +++ b/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 diff --git a/python/mocpy/tests/test_stmoc.py b/python/mocpy/tests/test_stmoc.py index a4706e78..06c046a2 100644 --- a/python/mocpy/tests/test_stmoc.py +++ b/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 diff --git a/python/mocpy/tests/test_tmoc.py b/python/mocpy/tests/test_tmoc.py index e727f2d0..a103cb93 100644 --- a/python/mocpy/tests/test_tmoc.py +++ b/python/mocpy/tests/test_tmoc.py @@ -1,11 +1,9 @@ +import numpy as np import pytest - -from astropy.time import Time - from astropy.io import ascii -from ..tmoc import TimeMOC, times_to_microseconds, microseconds_to_times +from astropy.time import Time -import numpy as np +from ..tmoc import TimeMOC, microseconds_to_times, times_to_microseconds def test_time_to_microsec_1(): @@ -24,8 +22,6 @@ def test_time_to_microsec_2(): us2 = np.asarray(t.jd * 86400000000, dtype=np.uint64) jd1 = microseconds_to_times(us1) jd2 = microseconds_to_times(us2) - print(us1) - print(us2) assert (us1 == us2).all() assert (jd1.jd == jd2.jd).all() @@ -55,6 +51,17 @@ def test_to_depth61_ranges(): ).all() +def test_n_cells(): + assert TimeMOC.n_cells(0) == 2 + with pytest.raises( + ValueError, + match=f"The depth should be comprised between 0 and {TimeMOC.MAX_ORDER}*", + ): + TimeMOC.n_cells(-1) + TimeMOC.n_cells(TimeMOC.MAX_ORDER + 1) + assert TimeMOC.n_cells(10) == 2 * TimeMOC.n_cells(9) + + def test_empty_tmoc(): times = Time([], format="jd", scale="tdb") tmoc = TimeMOC.from_times(times) @@ -85,9 +92,6 @@ def test_simple_tmoc(): assert tmoc.total_duration.sec == 2 * 1e-6 assert tmoc.max_order == 61 - # tmoc.write('tmoc.txt', format='json', overwrite=True) - tmoc.save("tmoc.txt", format="json", overwrite=True) - def test_single_time_tmoc(): times = Time(2 / TimeMOC.DAY_MICRO_SEC, format="jd", scale="tdb") @@ -208,8 +212,6 @@ def test_add_remove_back_and_forth(): times, delta_t=TimeMOC.order_to_time_resolution(61), ) - - print(tmoc.to_string()) tmoc.add_neighbours().remove_neighbours() assert tmoc == tmoc_expected @@ -307,43 +309,3 @@ def test_intersection(a, b, expect): res = a.intersection(b) assert res == expect - - -# ------- TESTING IO features -------# -def test_tmoc_save_load_deser(): - """Test of IO features. - - 1. Serializations - 2. Load - 3. Save - """ - tmoc = TimeMOC.from_string("31/1 32/4 35/") - - # test to_string "json" - tmoc_json = tmoc.to_string("json") - tmoc_from_json = TimeMOC.from_string(tmoc_json, "json") - assert tmoc == tmoc_from_json - - # test to_string "ascii" - tmoc_ascii = tmoc.to_string("ascii") - tmoc_from_ascii = TimeMOC.from_string(tmoc_ascii, "ascii") - assert tmoc == tmoc_from_ascii - - # test load from an ascii file - tmoc_load_ascii = TimeMOC.load("resources/MOC2.0/tmoc.ascii.txt", "ascii") - assert tmoc == tmoc_load_ascii - - # test load from fits - tmoc_load_fits = TimeMOC.load("resources/MOC2.0/TMOC.fits", "fits") - assert tmoc == tmoc_load_fits - - # tests of save in the three formats - tmoc.save("resources/MOC2.0/tmoc.py.test.fits", format="fits", overwrite=True) - tmoc.save("resources/MOC2.0/tmoc.py.test.json", format="json", overwrite=True) - tmoc.save("resources/MOC2.0/tmoc.py.test.ascii", format="ascii", overwrite=True) - tmoc_saved_fits = TimeMOC.load("resources/MOC2.0/tmoc.py.test.fits", "fits") - assert tmoc == tmoc_saved_fits - tmoc_saved_json = TimeMOC.load("resources/MOC2.0/tmoc.py.test.json", "json") - assert tmoc == tmoc_saved_json - tmoc_saved_ascii = TimeMOC.load("resources/MOC2.0/tmoc.py.test.ascii", "ascii") - assert tmoc == tmoc_saved_ascii diff --git a/python/mocpy/tmoc/tmoc.py b/python/mocpy/tmoc/tmoc.py index 9170a2d9..ab1b3005 100644 --- a/python/mocpy/tmoc/tmoc.py +++ b/python/mocpy/tmoc/tmoc.py @@ -1,11 +1,10 @@ import warnings -import numpy as np +import numpy as np from astropy.time import Time, TimeDelta -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" @@ -96,6 +95,33 @@ def max_order(self): depth = mocpy.get_tmoc_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.tmoc.TimeMOC.MAX_ORDER` + + Returns + ------- + int + The number of cells at the given order + + Examples + -------- + >>> from mocpy import TimeMOC + >>> TimeMOC.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_tmoc(depth) + def to_time_ranges(self): """Return the time ranges this TimeMOC contains.""" return microseconds_to_times(mocpy.to_ranges(self._store_index)) @@ -596,8 +622,8 @@ def plot(self, title="TimeMoc", view=(None, None), figsize=(9.5, 5), **kwargs): all the observation time window is rendered). """ - from matplotlib.colors import LinearSegmentedColormap import matplotlib.pyplot as plt + from matplotlib.colors import LinearSegmentedColormap if self.empty(): import warnings @@ -615,10 +641,7 @@ def plot(self, title="TimeMoc", view=(None, None), figsize=(9.5, 5), **kwargs): if max_jd < min_jd: raise ValueError( - "Invalid selection: max_jd = {} must be > to min_jd = {}".format( - max_jd, - min_jd, - ), + f"Invalid selection: max_jd = {max_jd} must be > to min_jd = {min_jd}", ) fig1 = plt.figure(figsize=figsize) diff --git a/src/lib.rs b/src/lib.rs index 293ecf5f..05dbb0d6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,7 +12,7 @@ use pyo3::{ }; use moc::{ - qty::{Frequency, Hpx, MocQty}, + qty::{Frequency, Hpx, MocQty, Time}, storage::u64idx::U64MocStore, utils, }; @@ -2019,6 +2019,36 @@ fn mocpy(_py: Python, m: &PyModule) -> PyResult<()> { .map_err(PyIOError::new_err) } + /// Gives the number of cells for a specific depth in a spatial MOC + /// + /// # Arguments + /// + /// * ``depth`` - The depth for which we want to know the corresponding number of healpix cells. + #[pyfn(m)] + fn n_cells_smoc(depth: u8) -> u64 { + Hpx::::n_cells(depth) + } + + /// Gives the number of cells for a specific depth in a time MOC + /// + /// # Arguments + /// + /// * ``depth`` - The depth for which we want to know the corresponding number of cells. + #[pyfn(m)] + fn n_cells_tmoc(depth: u8) -> u64 { + Time::::n_cells(depth) + } + + /// Gives the number of cells for a specific depth in a frequency MOC + /// + /// # Arguments + /// + /// * ``depth`` - The depth for which we want to know the corresponding number of cells. + #[pyfn(m)] + fn n_cells_fmoc(depth: u8) -> u64 { + Frequency::::n_cells(depth) + } + /// Create a spatial coverage from a list of HEALPix cell indices. /// /// # Arguments