create_enkf_experiment.py

import os.path
import argparse

import utilities.generic

import models.nemo.utilities.generic
import models.nemo.utilities.job_script
import models.nemo.utilities.argument_parser
import models.nemo.utilities.global_information
import models.nemo.utilities.base_experiment_plan
import models.nemo.utilities.configuration_namelist

import models.nemo.enkf_tools.utilities.generic
import models.nemo.enkf_tools.utilities.enkf_experiment
import models.nemo.enkf_tools.utilities.enkf_configuration
import models.nemo.enkf_tools.utilities.enkf_experiment_plan


#############
# functions #
#############

def create_enkf_experiment(version, configuration, name, start_date, assimilation_start_date,
  end_date, restart_frequency=models.nemo.utilities.generic.DEFAULT_VALUE_RESTART_FREQUENCY,
  assimilation_frequency=models.nemo.utilities.generic.DEFAULT_VALUE_RESTART_FREQUENCY,
  start_from=models.nemo.utilities.base_experiment_plan.START_FROM_REST,
  restart_name=None, restart_date=None, rebuilt_restart=False):
  '''
  Create a new EnKF experiment.
  '''

  print('Creating EnKF experiment "%s", in configuration "%s", in version "%s"...\n'%(name,
    configuration, version))

  # read the global information
  global_information = models.nemo.utilities.global_information.GlobalInformation()
  machine = global_information.get_machine()
  user_email = global_information.get_user_email()
  account = global_information.get_account()

  # create the experiment object
  experiment = models.nemo.enkf_tools.utilities.enkf_experiment.EnKFExperiment(version=version,
    configuration=configuration, name=name)
  configuration_path = experiment.get_configuration_path()
  experiment_path = experiment.create_experiment_directory()

  # read information in the configuration namelist
  configuration_namelist = models.nemo.utilities.configuration_namelist.ConfigurationNamelist(
      configuration_path=configuration_path, machine=machine)

  enkf_partition = configuration_namelist.get_enkf_partition()
  enkf_job_time = configuration_namelist.get_enkf_job_time()
  enkf_memory = configuration_namelist.get_enkf_memory()

  # copy files in the experiment directory
  template_directory_path = os.path.join(utilities.generic.get_toolbox_root_directory(), 'models',
    'nemo', 'enkf_tools', 'templates')

  enkf_namelist_path = os.path.join(template_directory_path,
    models.nemo.enkf_tools.utilities.generic.ENKF_NAMELIST_FILE_NAME)
  enkf_configuration_path = os.path.join(template_directory_path,
    models.nemo.enkf_tools.utilities.enkf_configuration.ENKF_CONFIGURATION_FILE_NAME)

  for file_path in (enkf_namelist_path, enkf_configuration_path):
    copied_file_path = os.path.join(experiment_path, os.path.basename(file_path))
    utilities.generic.copy_file(file_path=file_path, copied_file_path=copied_file_path)

  # prepare the assimilation script substitutions
  assimilation_script_substitutions = {}

  # - header
  assimilation_script_substitutions['#ASSIMILATION_SCRIPT_HEADER#'] = (
    models.nemo.utilities.job_script.create_script_header(machine=machine,
    partition=enkf_partition, job_name='enkf', job_time=enkf_job_time, exclusivity=False,
    number_of_cores=1, user_email=user_email, account=account, memory=enkf_memory))

  # - paths
  assimilation_script_substitutions['#CLIMATETOOLBOX_PATH#'] = (
    utilities.generic.get_toolbox_root_directory())

  # - experiment information
  assimilation_script_substitutions['#VERSION#'] = version
  assimilation_script_substitutions['#CONFIGURATION#'] = configuration
  assimilation_script_substitutions['#EXPERIMENT_NAME#'] = name

  # copy the assimilation script in the experiment directory (with adjustments)
  assimilation_script_path = os.path.join(template_directory_path,
    models.nemo.utilities.job_script.ASSIMILATION_SCRIPT_FILE_NAME)
  updated_assimilation_script_path = os.path.join(experiment_path,
    models.nemo.utilities.job_script.ASSIMILATION_SCRIPT_FILE_NAME)
  assimilation_script = models.nemo.utilities.job_script.JobScript(
    file_path=assimilation_script_path)
  assimilation_script.update(substitutions=assimilation_script_substitutions,
    path=updated_assimilation_script_path)

  # create the experiment plan
  experiment_plan = models.nemo.enkf_tools.utilities.enkf_experiment_plan.EnKFExperimentPlan(
    experiment=experiment)
  experiment_plan.generate_legs(start_date=start_date,
    assimilation_start_date=assimilation_start_date, end_date=end_date,
    restart_frequency=restart_frequency, assimilation_frequency=assimilation_frequency)
  experiment_plan.set_start_information(start_from=start_from, restart_name=restart_name,
    restart_date=restart_date, rebuilt_restart=rebuilt_restart)
  experiment_plan.write()


#################
# main function #
#################

def main():

  print()

  parser = argparse.ArgumentParser(description=create_enkf_experiment.__doc__,
    formatter_class=argparse.RawTextHelpFormatter)
  models.nemo.utilities.argument_parser.add_argument_version(parser=parser, required=False)
  models.nemo.utilities.argument_parser.add_argument_configuration(parser=parser, required=False)
  models.nemo.utilities.argument_parser.add_argument_name(parser=parser)
  models.nemo.utilities.argument_parser.add_argument_start_date(parser=parser, help_choice='enkf')
  parser.add_argument(
    '-asd', '--assimilation_start_date',
    required=True,
    help=('start date of the assimilated part of the experiment'
      + '\n(format: %s)'%models.nemo.utilities.generic.DATE_FORMAT),
  )
  models.nemo.utilities.argument_parser.add_argument_end_date(parser=parser)
  models.nemo.utilities.argument_parser.add_argument_restart_frequency(parser=parser,
    help_choice='enkf')
  parser.add_argument(
    '-af', '--assimilation_frequency',
    default=models.nemo.utilities.generic.DEFAULT_VALUE_RESTART_FREQUENCY,
    help=('frequency at which data assimilation occurs (and restarts are created) during the '
      + 'assimilated part of the experiment'
      + '\n(format: %s)'%models.nemo.utilities.generic.RESTART_FREQUENCY_FORMAT
      + '\n(default: %s)'%models.nemo.utilities.generic.DEFAULT_VALUE_RESTART_FREQUENCY),
  )
  models.nemo.utilities.argument_parser.add_argument_start_from(parser=parser)
  models.nemo.utilities.argument_parser.add_argument_restart_name(parser=parser)
  models.nemo.utilities.argument_parser.add_argument_restart_date(parser=parser)
  models.nemo.utilities.argument_parser.add_argument_rebuilt_restart(parser=parser)
  arguments = vars(parser.parse_args())

  try:
    arguments['version'], arguments['configuration'] = (
      models.nemo.utilities.global_information.get_version_and_configuration(
      tentative_version=arguments['version'], tentative_configuration=arguments['configuration']))
  except utilities.generic.ClimateToolboxException as e:
    e.display()
    return

  try:
    create_enkf_experiment(**arguments)
  except utilities.generic.ClimateToolboxException as e:
    e.display()
    return


if __name__ == '__main__':
  main()