pmsco-public/pmsco/database/util.py

import logging
import numpy as np
from pathlib import Path
import pmsco.dispatch as dispatch

logger = logging.getLogger(__name__)


## mapping of database fields to special parameter names
#
# `_db` parameters are returned by some query methods to identify the database records.
#
DB_SPECIAL_PARAMS = {"project_id": "_db_project_id",
                     "job_id": "_db_job_id",
                     "model_id": "_db_model_id",
                     "result_id": "_db_result_id",
                     "model": "_model",
                     "scan": "_scan",
                     "domain": "_domain",
                     "emit": "_emit",
                     "region": "_region",
                     "gen": "_gen",
                     "particle": "_particle",
                     "rfac": "_rfac",
                     "secs": "_secs",
                     "timestamp": "_timestamp"}


## numpy data types of special parameters by database field
#
# this dictionary helps to create a numpy array from a database record.
#
DB_SPECIAL_NUMPY_TYPES = {"_db_project_id": "i8",
                          "_db_job_id": "i8",
                          "_db_model_id": "i8",
                          "_db_result_id": "i8",
                          "_model": "i8",
                          "_scan": "i8",
                          "_domain": "i8",
                          "_emit": "i8",
                          "_region": "i8",
                          "_gen": "i8",
                          "_particle": "i8",
                          "_rfac": "f8",
                          "_secs": "f8",
                          "_timestamp": "f8"}


def regular_params(d):
    """
    filter regular parameters from dictionary

    returns a dictionary containing only the regular parameters (those not prefixed with an underscore).

    @param d: dict or numpy.void or pmsco.dispatch.CalcID.
        the param names must have no leading underscore.
        the numpy.void type occurs when an element of a structured array is extracted.
        the CalcID does not contain a regular parameter and will return an empty dictionary.
        it is supported only for compatibility with special_params function.
        a tuple or list is interpreted as a sequence of parameter names.
        in this case the names representing special parameters are returned with underscore removed.

    @return: dict for mapping types (numpy and dict) containing the regular key: value pairs of the original object.
        list (tuple) of parameter names for sequence (tuple) types.
        leading underscores are removed from key names.
    """
    if isinstance(d, np.void):
        d = {k: d[k] for k in d.dtype.names if k[0] != "_"}
    elif isinstance(d, dispatch.CalcID):
        d = {}
    elif isinstance(d, tuple):
        d = [k for k in d if k[0] != "_"]
        d = tuple(d)
    elif isinstance(d, dict):
        d = {k: v for k, v in d.items() if k[0] != "_"}
    else:
        d = [k for k in d if k[0] != "_"]

    return d


def special_params(d):
    """
    filter special parameters from model dictionary, numpy record or sequence.

    special parameters are those prefixed with an underscore.
    the underscore is removed from the keys.
    fields starting with '_db_' are removed.

    @param d: dict or numpy.void or pmsco.dispatch.CalcID or sequence.
        in the case of a dict or numpy.void,
        the key names of the special parameters must have a leading underscore.
        the numpy.void type occurs when an element of a structured array is extracted.
        in the case of a CalcID, the attribute names become the key names.
        a tuple or list is interpreted as a sequence of parameter names.
        in this case the names representing special parameters are returned with underscore removed.

    @return
        the return type depends on the type of input `d`:
        @arg in the case of a dict, numpy.void or CalcID it is a dictionary.
        @arg in the case of a tuple or list the return type is the same as the input.
    """
    if isinstance(d, np.void):
        d = {k[1:]: d[k] for k in d.dtype.names if k[0] == "_" and k[0:4] != "_db_"}
    elif isinstance(d, dispatch.CalcID):
        d = d._asdict()
    elif isinstance(d, tuple):
        d = [k[1:] for k in d if k[0] == "_" and k[0:4] != "_db_"]
        d = tuple(d)
    elif isinstance(d, dict):
        d = {k[1:]: v for k, v in d.items() if k[0] == "_" and k[0:4] != "_db_"}
    else:
        d = [k[1:] for k in d if k[0] == "_" and k[0:4] != "_db_"]

    return d


def field_to_param(f):
    """
    translate database field name to parameter name.

    field names of optimization parameters are unchanged.
    special parameters are prefixed by '_' or '_db_'.

    @param f: (str) database field name.
    @return: (str) parameter name as used in model dictionaries.
    """
    try:
        p = DB_SPECIAL_PARAMS[f]
    except KeyError:
        p = f
    return p


def field_to_numpy_type(f):
    """
    determine the numpy data type string of a database field.

    @param f: (str) database field name.
    @return: (str) numpy type description, e.g. 'f8'.
    """
    try:
        t = DB_SPECIAL_NUMPY_TYPES[f]
    except KeyError:
        t = 'f8'
    return t


def is_sqlite3_file(path_like):
    """
    test whether a file is an sqlite3 database file.

    @param path_like: file path (str or pathlib.Path).
    @return: (bool)
    """
    try:
        with Path(path_like).open("rb") as f:
            s = f.read(16)
            return s == b"SQLite format 3\000"
    except OSError:
        return False