Source code for nowcast.workers.make_boundary_conditions

# Copyright 2016-2019 Doug Latornell, 43ravens
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""GoMSS NEMO nowcast worker to run Matlab scripts that generate daily
boundary conditions files from downloaded Mercator Ocean PSY4 model fields.
"""
import glob
import logging
import shlex
import subprocess
from pathlib import Path

import arrow
import numpy
import xarray

from nemo_nowcast import NowcastWorker, WorkerError

NAME = "make_boundary_conditions"
logger = logging.getLogger(NAME)


def main():
    """Set up and run the worker.

    For command-line usage see:

    :command:`python -m nowcast.workers.make_boundary_conditions --help`
    """
    worker = NowcastWorker(NAME, description=__doc__)
    worker.init_cli()
    worker.cli.add_date_option(
        "--psy4-date",
        default=arrow.now().floor("day"),
        help="Date for which to generate boundary conditions files.",
    )
    worker.run(make_boundary_conditions, success, failure)


def success(parsed_args):
    ymd = parsed_args.psy4_date.format("YYYY-MM-DD")
    logger.info(
        f"{ymd} boundary conditions files generation complete", extra={"psy4_date": ymd}
    )
    msg_type = "success"
    return msg_type


def failure(parsed_args):
    ymd = parsed_args.psy4_date.format("YYYY-MM-DD")
    logger.critical(
        f"{ymd} boundary conditions files generation failed", extra={"psy4_date": ymd}
    )
    msg_type = "failure"
    return msg_type


def make_boundary_conditions(parsed_args, config, *args):
    start_ymd = parsed_args.psy4_date.format("YYYY-MM-DD")
    logger.info(
        f"generating open boundary conditions files for {start_ymd}",
        extra={"psy4_date": start_ymd},
    )
    bdy_config = config["open boundaries"]
    forecast_days = bdy_config["download"]["forecast days"]
    end_ymd = parsed_args.psy4_date.shift(days=forecast_days - 1).format("YYYY-MM-DD")
    bcs_config = bdy_config["boundary conditions"]
    matlab_scripts = bcs_config["matlab scripts"]
    dest_dir = Path(bcs_config["dest dir"])
    filename_tmpl = bcs_config["file template"]
    cmd = shlex.split("matlab -nodesktop -nodisplay -nosplash -r ")
    cmd.append(
        f"cd(fileparts('{matlab_scripts}')); "
        f"gen_bdy_loop_4D('{start_ymd}', '{end_ymd}'); "
        f"exit;"
    )
    logger.debug(f"running command in subprocess: {cmd}")
    proc = subprocess.run(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True
    )
    for line in proc.stdout.splitlines():
        line = line.strip()
        if line:
            logger.debug(line)
    files = []
    start = parsed_args.psy4_date
    end = parsed_args.psy4_date.shift(days=forecast_days - 1)
    for day in arrow.Arrow.range("day", start, end):
        file_pattern = filename_tmpl.format(
            bdy="*",
            var="*",
            yyyy=day.format("YYYY"),
            mm=day.format("MM"),
            dd=day.format("DD"),
        )
        files.extend(glob.glob(str(dest_dir / file_pattern)))
    for file in files:
        with xarray.open_dataset(file) as ds:
            if any(numpy.any(ds.variables[var].isnull()) for var in ds.variables):
                logger.critical(f"found None or NaN value(s) in {file}")
                raise WorkerError
    checklist = {"psy4 date": start_ymd, "files": files}
    return checklist


if __name__ == "__main__":
    main()  # pragma: no cover