Merge pull request #58 from sot/new-utils-logger-cm-random-ra-dec

New utils: logger context manager and random RA and Decs

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__)

ska_helpers/chandra_models.py

@@ -2,6 +2,7 @@
 """
 Get data from chandra_models repository.
 """
+
 import contextlib
 import functools
 import hashlib

ska_helpers/git_helpers.py

@@ -2,6 +2,7 @@
 """
 Helper functions for using git.
 """
+
 import functools
 import git
 import re

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

ska_helpers/tests/test_utils.py

@@ -1,18 +1,24 @@
 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,
     LRUDict,
     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)

ska_helpers/utils.py

@@ -5,17 +5,49 @@
 import os
 from collections import OrderedDict
 
+import numpy as np
+
 __all__ = [
+    "get_owner",
     "LazyDict",
     "LazyVal",
     "LRUDict",
     "lru_cache_timed",
     "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.