diff --git a/ska_helpers/__init__.py b/ska_helpers/__init__.py index 2511c97..35319a2 100644 --- a/ska_helpers/__init__.py +++ b/ska_helpers/__init__.py @@ -2,6 +2,7 @@ """ Ska_helpers is a collection of utilities for the Ska3 runtime environment. """ + from .version import get_version __version__ = get_version(__package__) diff --git a/ska_helpers/chandra_models.py b/ska_helpers/chandra_models.py index 4abcfa8..13fbe80 100644 --- a/ska_helpers/chandra_models.py +++ b/ska_helpers/chandra_models.py @@ -2,6 +2,7 @@ """ Get data from chandra_models repository. """ + import contextlib import functools import hashlib diff --git a/ska_helpers/git_helpers.py b/ska_helpers/git_helpers.py index a13eee6..166e6f5 100644 --- a/ska_helpers/git_helpers.py +++ b/ska_helpers/git_helpers.py @@ -2,6 +2,7 @@ """ Helper functions for using git. """ + import functools import git import re diff --git a/ska_helpers/retry/__init__.py b/ska_helpers/retry/__init__.py index da44ce6..275cf6b 100644 --- a/ska_helpers/retry/__init__.py +++ b/ska_helpers/retry/__init__.py @@ -19,6 +19,7 @@ See the License for the specific language governing permissions and limitations under the License. """ + __all__ = ["retry", "retry_call", "RetryError", "tables_open_file"] import logging diff --git a/ska_helpers/tests/test_utils.py b/ska_helpers/tests/test_utils.py index 56e54cd..816c51d 100644 --- a/ska_helpers/tests/test_utils.py +++ b/ska_helpers/tests/test_utils.py @@ -1,11 +1,15 @@ import functools +import logging import os import pickle import time from dataclasses import dataclass +import agasc +import numpy as np import pytest +import ska_helpers.logging from ska_helpers.utils import ( LazyDict, LazyVal, @@ -13,6 +17,8 @@ TypedDescriptor, convert_to_int_float_str, lru_cache_timed, + random_radec_in_cone, + set_log_level, temp_env_var, ) @@ -223,3 +229,43 @@ def test_int_descriptor_is_required_has_default_exception(cls_descriptor): @dataclass class MyClass: quat: int = cls_descriptor(default=30, required=True) + + +def test_set_log_level(): + logger = ska_helpers.logging.basic_logger("test_utils", level="DEBUG") + + assert logger.level == logging.DEBUG + assert len(logger.handlers) == 1 + for hdlr in logger.handlers: + assert hdlr.level == 0 + + with set_log_level(logger, "INFO"): + assert logger.level == logging.INFO + assert len(logger.handlers) == 1 + for hdlr in logger.handlers: + assert hdlr.level == logging.INFO + + assert logger.level == logging.DEBUG + assert len(logger.handlers) == 1 + for hdlr in logger.handlers: + assert hdlr.level == 0 + + +def test_random_radec_in_cone_scalar(): + np.random.seed(0) + ra, dec = random_radec_in_cone(10, 20, angle=5) + assert np.isclose(ra, 8.6733489) + assert np.isclose(dec, 15.964518) + + +def test_random_radec_in_cone_size_values(): + np.random.seed(0) + ra, dec = random_radec_in_cone(10, 20, angle=5, size=2) + assert np.allclose(ra, [8.77992603, 6.18623754]) + assert np.allclose(dec, [16.29571322, 19.15880785]) + + +def test_random_radec_in_cone_size_angle(): + np.random.seed(0) + ra, dec = random_radec_in_cone(10, 20, angle=5, size=10000) + assert np.all(agasc.sphere_dist(ra, dec, 10, 20) < 5) diff --git a/ska_helpers/utils.py b/ska_helpers/utils.py index e3a4aee..702af74 100644 --- a/ska_helpers/utils.py +++ b/ska_helpers/utils.py @@ -5,7 +5,10 @@ import os from collections import OrderedDict +import numpy as np + __all__ = [ + "get_owner", "LazyDict", "LazyVal", "LRUDict", @@ -13,9 +16,38 @@ "temp_env_var", "convert_to_int_float_str", "TypedDescriptor", + "set_log_level", + "random_radec_in_cone", ] +@contextlib.contextmanager +def set_log_level(logger, level=None): + """Set the log level of a logger and its handlers for context block. + + Parameters + ---------- + logger : logging.Logger + The logger object to set the level for. + level : str, int, None, optional + The log level to set. This can be a string like "DEBUG", "INFO", "WARNING", + "ERROR", "CRITICAL", or an integer value from the ``logging`` module. If level + is None (default), the log level is not changed. + """ + orig_levels = {} + if level is not None: + orig_levels[logger] = logger.level + logger.setLevel(level) + for handler in logger.handlers: + orig_levels[handler] = handler.level + handler.setLevel(level) + try: + yield + finally: + for log_obj, orig_level in orig_levels.items(): + log_obj.setLevel(orig_level) + + def get_owner(path): """ Returns the owner of a file or directory. @@ -244,6 +276,62 @@ def cache_info(): return _wrapper +def random_radec_in_cone( + ra: float, dec: float, *, angle: float, size=None +) -> tuple[np.ndarray, np.ndarray]: + """Get random sky coordinates within a cone. + + This returns a tuple of RA and Dec values within ``angle`` degrees of ``ra`` and + ``dec``. The coordinates are uniformly distributed over the sky area. + + Parameters + ---------- + ra : float + RA in degrees of the center of the cone. + dec : float + Dec in degrees of the center of the cone. + angle : float + The radius of the cone in degrees. + size : int, optional + The number of random coordinates to generate. If not specified, a single + coordinate is generated. + + Returns + ------- + ra_rand : np.ndarray + Random RA values in degrees. + dec_rand : np.ndarray + Random Dec values in degrees. + """ + import chandra_aca.transform as cat + from Quaternion import Quat + + # Convert input angles from degrees to radians + angle_rad = np.radians(angle) + + # Generate a random azimuthal angle (phi) between 0 and 2π + phi = np.random.uniform(0, 2 * np.pi, size=size) + + # Generate a random polar angle (theta) within the specified angle from the north pole + u = np.random.uniform(0, 1, size=size) + theta = np.arccos(1 - u * (1 - np.cos(angle_rad))) + + # Generate vectors around pole (dec=90) + ra_rot = np.degrees(phi) + dec_rot = 90 - np.degrees(theta) + eci = cat.radec_to_eci(ra_rot, dec_rot) + + # Swap x and z axes to get vectors centered around RA=0 ad Dec=0 + eci[..., [0, 2]] = eci[..., [2, 0]] + + # Now rotate the random vectors to be centered about the desired RA and Dec. + q = Quat([ra, dec, 0]) + eci_rot = q.transform @ eci.T + ra_rand, dec_rand = cat.eci_to_radec(eci_rot.T) + + return ra_rand, dec_rand + + class LRUDict(OrderedDict): """ Dict that maintains a fixed capacity and evicts least recently used item when full.