Source code for compass.ocean.tests.global_ocean.metadata

import os
import shutil
import subprocess
from datetime import datetime

import numpy
import xarray


def get_author_and_email_from_git(config):
    """
    Get the author and email address config options from git if they are set to
    autodetect

    Parameters
    ----------
    config : compass.config.CompassConfigParser
        Configuration options for this test case
    """
    author = config.get('global_ocean', 'author')
    if author == 'autodetect':
        try:
            author = subprocess.check_output(
                ['git', 'config', 'user.name']).decode("utf-8").strip()
        except subprocess.CalledProcessError:
            raise ValueError('It appears you have not set up a git username '
                             'yet.  Please do so or provide a config file '
                             'that sets config option "author" in '
                             '[global_ocean].')
        config.set('global_ocean', 'author', author)

    email = config.get('global_ocean', 'email')
    if email == 'autodetect':
        try:
            email = subprocess.check_output(
                ['git', 'config', 'user.email']).decode("utf-8").strip()
        except subprocess.CalledProcessError:
            raise ValueError('It appears you have not set your email in git '
                             'yet.  Please do so or provide a config file '
                             'that sets config option "email" in '
                             '[global_ocean].')
        config.set('global_ocean', 'email', email)




def get_e3sm_mesh_names(config, levels):
    """
    Get short and long E3SM mesh name from config options and the given number
    of vertical levels (typically taken from an initial condition or restart
    file).

    Parameters
    ----------
    config : compass.config.CompassConfigParser
        Configuration options for this test case

    levels : int
        The number of vertical levels

    Returns
    -------
    short_mesh_name : str
        The short E3SM name of the ocean and sea-ice mesh

    long_mesh_name : str
        The long E3SM name of the ocean and sea-ice mesh
    """

    config.set('global_ocean', 'levels', f'{levels}')
    mesh_prefix = config.get('global_ocean', 'prefix')
    min_res = config.get('global_ocean', 'min_res')
    max_res = config.get('global_ocean', 'max_res')
    e3sm_version = config.get('global_ocean', 'e3sm_version')
    mesh_revision = config.get('global_ocean', 'mesh_revision')

    if min_res == max_res:
        res = min_res
    else:
        res = f'{min_res}to{max_res}'

    short_mesh_name = f'{mesh_prefix}{res}E{e3sm_version}r{mesh_revision}'
    long_mesh_name = \
        f'{mesh_prefix}{res}kmL{levels}E3SMv{e3sm_version}r{mesh_revision}'

    return short_mesh_name, long_mesh_name





def add_mesh_and_init_metadata(output_filenames, config, init_filename):
    """
    Add MPAS mesh and initial condition metadata to NetCDF outputs of the given
    step

    Parameters
    ----------
    output_filenames : list
        A list of output files.

    config : compass.config.CompassConfigParser
        Configuration options for this test case

    init_filename : str
        The name of an initial condition file to get the number of vertical
        levels and maximum depth from
    """

    if config.getboolean('global_ocean', 'add_metadata'):
        with xarray.open_dataset(init_filename) as dsInit:
            metadata = _get_metadata(dsInit, config)

        for filename in output_filenames:
            if filename.endswith('.nc'):
                args = ['ncks']
                for key, value in metadata.items():
                    args.extend(['--glb_att_add', f'{key}={value}'])
                name, ext = os.path.splitext(filename)
                new_filename = f'{name}_with_metadata{ext}'
                args.extend([filename, new_filename])
                subprocess.check_call(args)
                shutil.move(new_filename, filename)



def _get_metadata(dsInit, config):
    """ add metadata to a given dataset """

    author = config.get('global_ocean', 'author')
    email = config.get('global_ocean', 'email')
    creation_date = config.get('global_ocean', 'creation_date')
    if creation_date == 'autodetect':
        now = datetime.now()
        creation_date = now.strftime("%Y%m%d")
        config.set('global_ocean', 'creation_date', creation_date)

    max_depth = dsInit.bottomDepth.max().values
    # round to the nearest 0.1 m
    max_depth = numpy.round(max_depth, 1)
    config.set('global_ocean', 'max_depth', f'{max_depth}')

    mesh_prefix = config.get('global_ocean', 'prefix')
    min_res = config.get('global_ocean', 'min_res')
    max_res = config.get('global_ocean', 'max_res')
    levels = dsInit.sizes['nVertLevels']
    config.set('global_ocean', 'levels', f'{levels}')
    e3sm_version = config.get('global_ocean', 'e3sm_version')
    mesh_revision = config.get('global_ocean', 'mesh_revision')
    pull_request = config.get('global_ocean', 'pull_request')

    short_name, long_name = get_e3sm_mesh_names(config, levels)

    descriptions = dict()

    replacements = {'<<<levels>>>': f'{levels}',
                    '<<<max_depth>>>': f'{max_depth:g}',
                    '<<<creation_date>>>': creation_date}

    for prefix in ['mesh', 'init', 'bathy', 'bgc', 'wisc']:
        option = f'{prefix}_description'
        if config.has_option('global_ocean', option):
            description = config.get('global_ocean', option)
            description = ' '.join(
                [line.strip() for line in description.split('\n')])
            for placeholder, replacement in replacements.items():
                if placeholder in description:
                    description = description.replace(placeholder,
                                                      replacement)
            descriptions[prefix] = description

    prefix = f'MPAS_Mesh_{mesh_prefix}'

    metadata = {'MPAS_Mesh_Short_Name': short_name,
                'MPAS_Mesh_Long_Name': long_name,
                'MPAS_Mesh_Prefix': mesh_prefix,
                'MPAS_Mesh_E3SM_Version': e3sm_version,
                'MPAS_Mesh_Pull_Request': pull_request,
                f'{prefix}_Revision': mesh_revision,
                f'{prefix}_Version_Author': author,
                f'{prefix}_Version_Author_Email': email,
                f'{prefix}_Version_Creation_Date': creation_date,
                f'{prefix}_Minimum_Resolution_km': min_res,
                f'{prefix}_Maximum_Resolution_km': max_res,
                f'{prefix}_Maximum_Depth_m': f'{max_depth}',
                f'{prefix}_Number_of_Levels': f'{levels}',
                'MPAS_Mesh_Description': descriptions['mesh'],
                'MPAS_Mesh_Bathymetry': descriptions['bathy'],
                'MPAS_Initial_Condition': descriptions['init']}

    if 'wisc' in descriptions:
        metadata['MPAS_Mesh_Ice_Shelf_Cavities'] = descriptions['wisc']

    if 'bgc' in descriptions:
        metadata['MPAS_Mesh_Biogeochemistry'] = descriptions['bgc']

    packages = {'compass': 'compass', 'JIGSAW': 'jigsaw',
                'JIGSAW_Python': 'jigsawpy', 'MPAS_Tools': 'mpas_tools',
                'NCO': 'nco', 'ESMF': 'esmf',
                'geometric_features': 'geometric_features',
                'Metis': 'metis', 'pyremap': 'pyremap'}

    for name in packages:
        package = packages[name]
        metadata[f'MPAS_Mesh_{name}_Version'] = \
            _get_conda_package_version(package)

    return metadata


def _get_conda_package_version(package):
    conda = subprocess.check_output(['conda', 'list', package]).decode("utf-8")
    lines = conda.split('\n')
    for line in lines:
        parts = line.split()
        if len(parts) > 0 and parts[0] == package:
            return parts[1]

    return 'not found'