Implement logger for AbacusHOD

abacusnbody/hod/__init__.py

@@ -0,0 +1,3 @@
+from .utils import setup_logging
+
+__all__ = ['setup_logging']

abacusnbody/hod/abacus_hod.py

@@ -9,7 +9,7 @@
 import gc
 import time
 from pathlib import Path
-import warnings
+import logging
 
 import asdf
 import h5py
@@ -99,6 +99,7 @@ def __init__(
         n_chunks: int, optional
             Number of chunks to split the input from the halo+particle subsample and number of output files in which to write out the galaxy catalogs following the format ``{tracer}s_{chunk}.dat``.
         """
+        self.logger = logging.getLogger('AbacusHOD')
         # simulation details
         self.sim_name = sim_params['sim_name']
         self.sim_dir = sim_params['sim_dir']
@@ -385,7 +386,7 @@ def staging(self):
         halo_ticker = 0
         parts_ticker = 0
         for eslab in range(start, end):
-            print('Loading simulation by slab, ', eslab)
+            self.logger.info(f'Loading simulation slab {eslab}')
             if (
                 ('ELG' not in self.tracers.keys())
                 and ('QSO' not in self.tracers.keys())
@@ -427,7 +428,7 @@ def staging(self):
                 ]  # halo velocity dispersions, km/s
 
             if len(halo_vel_dev.shape) == 1:
-                warnings.warn(
+                self.logger.warning(
                     'Warning: galaxy x, y velocity bias randoms not set, using z randoms instead. x, y velocities may be unreliable.'
                 )
                 halo_vel_dev = np.concatenate(
@@ -555,7 +556,7 @@ def staging(self):
 
         # sort halos by hid, important for conformity
         if not np.all(hid[:-1] <= hid[1:]):
-            print('sorting halos for conformity calculation')
+            self.logger.info('Sorting halos for conformity calculation.')
             sortind = np.argsort(hid)
             hpos = hpos[sortind]
             hvel = hvel[sortind]
@@ -750,7 +751,9 @@ def run_hod(
                 )
             self.particle_data['prandoms'] = r3
 
-            print('gen randoms took, ', time.time() - start)
+            self.logger.info(
+                f'Randoms generated in elapsed time {time.time() - start:.2f} s.'
+            )
 
         start = time.time()
         mock_dict = gen_gal_cat(
@@ -768,7 +771,7 @@ def run_hod(
             verbose=verbose,
             fn_ext=fn_ext,
         )
-        print('gen mocks', time.time() - start)
+        self.logger.info(f'HOD generated in elapsed time {time.time() - start:.2f} s.')
 
         return mock_dict
 

abacusnbody/hod/utils.py

@@ -0,0 +1,120 @@
+# BSD 3-Clause License
+
+# Copyright (c) 2021, cosmodesi
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+# taken from https://github.com/cosmodesi/desilike/blob/main/desilike/utils.py
+
+import logging
+import sys
+import os
+import time
+import traceback
+
+
+def setup_logging(
+    level=logging.INFO, stream=sys.stdout, filename=None, filemode='w', **kwargs
+):
+    """
+    Set up logging.
+
+    Parameters
+    ----------
+    level : string, int, default=logging.INFO
+        Logging level.
+
+    stream : _io.TextIOWrapper, default=sys.stdout
+        Where to stream.
+
+    filename : string, default=None
+        If not ``None`` stream to file name.
+
+    filemode : string, default='w'
+        Mode to open file, only used if filename is not ``None``.
+
+    kwargs : dict
+        Other arguments for :func:`logging.basicConfig`.
+    """
+    # Cannot provide stream and filename kwargs at the same time to logging.basicConfig, so handle different cases
+    # Thanks to https://stackoverflow.com/questions/30861524/logging-basicconfig-not-creating-log-file-when-i-run-in-pycharm
+    if isinstance(level, str):
+        level = {
+            'info': logging.INFO,
+            'debug': logging.DEBUG,
+            'warning': logging.WARNING,
+        }[level.lower()]
+    for handler in logging.root.handlers:
+        logging.root.removeHandler(handler)
+
+    t0 = time.time()
+
+    class MyFormatter(logging.Formatter):
+        def format(self, record):
+            self._style._fmt = (
+                '[%09.2f] ' % (time.time() - t0)
+                + ' %(asctime)s %(name)-28s %(levelname)-8s %(message)s'
+            )
+            return super(MyFormatter, self).format(record)
+
+    fmt = MyFormatter(datefmt='%m-%d %H:%M ')
+    if filename is not None:
+        mkdir(os.path.dirname(filename))
+        handler = logging.FileHandler(filename, mode=filemode)
+    else:
+        handler = logging.StreamHandler(stream=stream)
+    handler.setFormatter(fmt)
+    logging.basicConfig(level=level, handlers=[handler], **kwargs)
+    sys.excepthook = exception_handler
+
+
+def exception_handler(exc_type, exc_value, exc_traceback):
+    """Print exception with a logger."""
+    # Do not print traceback if the exception has been handled and logged
+    _logger_name = 'Exception'
+    log = logging.getLogger(_logger_name)
+    line = '=' * 100
+    # log.critical(line[len(_logger_name) + 5:] + '\n' + ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback)) + line)
+    log.critical(
+        '\n'
+        + line
+        + '\n'
+        + ''.join(traceback.format_exception(exc_type, exc_value, exc_traceback))
+        + line
+    )
+    if exc_type is KeyboardInterrupt:
+        log.critical('Interrupted by the user.')
+    else:
+        log.critical('An error occured.')
+
+
+def mkdir(dirname):
+    """Try to create ``dirname`` and catch :class:`OSError`."""
+    try:
+        os.makedirs(dirname)  # MPI...
+    except OSError:
+        return