Skip to content
Extraits de code Groupes Projets
create_enkf_experiment.py 14,5 ko
Newer Older
  • Learn to ignore specific revisions
  • import models.nemo.utilities.argument_parser
    import models.nemo.utilities.global_information
    
    import models.nemo.utilities.base_experiment_plan
    
    Antoine Barthélemy's avatar
    Antoine Barthélemy a validé
    import models.nemo.utilities.configuration_namelist
    
    import models.nemo.enkf_tools.utilities.generic
    
    import models.nemo.enkf_tools.utilities.argument_parser
    
    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
    
    #############
    # variables #
    #############
    
    ASSIMILATION_SCRIPT_PYTHON_MODULE_FILES = {}
    
    
    
    #############
    # 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_ASSIMILATION_FREQUENCY,
    
      start_from=models.nemo.utilities.base_experiment_plan.START_FROM_REST,
    
      restart_name=None, restart_date=None, rebuilt_restart=False, restart_from_enkf=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)
    
    Antoine Barthélemy's avatar
    Antoine Barthélemy a validé
      configuration_path = experiment.get_configuration_path()
    
      experiment_path = experiment.create_experiment_directory()
    
      # read various parameters in the reference experiment ice namelist
    
      ocean_namelist_path = os.path.join(configuration_path,
        models.nemo.utilities.generic.REFERENCE_EXPERIMENT_NAME,
        models.nemo.utilities.generic.OCEAN_NAMELIST_FILE_NAME)
      ocean_namelist = models.nemo.utilities.namelist.Namelist(file_path=ocean_namelist_path)
    
      ocean_parameters = [
        ('rn_Dt', float),
        ('nn_fsbc', int),
      ]
      ocean_namelist_parameters = ocean_namelist.extract_parameters(parameters=ocean_parameters)
    
      ice_namelist_path = os.path.join(configuration_path,
    
        models.nemo.utilities.generic.REFERENCE_EXPERIMENT_NAME,
        models.nemo.utilities.generic.ICE_NAMELIST_FILE_NAME)
    
      ice_namelist = models.nemo.utilities.namelist.Namelist(file_path=ice_namelist_path)
    
        ('nlay_i', int),
        ('nlay_s', int),
        ('rn_amax_n', float),
        ('rn_amax_s', float),
    
        ('rn_himin', float),
        ('nn_icesal', int),
    
      ice_namelist_parameters = ice_namelist.extract_parameters(parameters=ice_parameters)
    
    Antoine Barthélemy's avatar
    Antoine Barthélemy a validé
      # 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()
    
    
      # prepare the module loading lines
      if machine in ASSIMILATION_SCRIPT_PYTHON_MODULE_FILES:
        toolbox_module_file = utilities.module_file.ModuleFile(
          file_path=os.path.join(utilities.generic.get_toolbox_root_directory(), 'machines',
          'modules', ASSIMILATION_SCRIPT_PYTHON_MODULE_FILES[machine]))
        toolbox_module_loading_lines = toolbox_module_file.extract_loading_lines()
      else:
        toolbox_module_loading_lines = 'echo "No modules needed"'
    
    
      enkf_module_file = utilities.module_file.ModuleFile(
        file_path=os.path.join(utilities.generic.get_toolbox_root_directory(), 'models', 'nemo',
        'enkf_tools', 'modules', (machine + '.txt')))
    
      enkf_module_loading_lines = enkf_module_file.extract_loading_lines(spaces_to_prepend='  ')
    
      # prepare the assimilation script substitutions
      assimilation_script_substitutions = {}
    
      # - header
      assimilation_script_substitutions['#ASSIMILATION_SCRIPT_HEADER#'] = (
    
    Antoine Barthélemy's avatar
    Antoine Barthélemy a validé
        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
    
    
      # - modules
      assimilation_script_substitutions['#CLIMATETOOLBOX_MODULE_LOADING_1#'] = (
        toolbox_module_loading_lines)
      assimilation_script_substitutions['#CLIMATETOOLBOX_MODULE_LOADING_2#'] = (
        toolbox_module_loading_lines)
    
      assimilation_script_substitutions['#ENKF_MODULE_LOADING#'] = (
        enkf_module_loading_lines)
    
      # copy the assimilation script in the experiment directory (with adjustments)
    
      template_directory_path = os.path.join(utilities.generic.get_toolbox_root_directory(), 'models',
        'nemo', 'enkf_tools', 'templates')
    
      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)
    
    
      # prepare the dummy assimilation script substitutions
      dummy_assimilation_script_substitutions = {}
    
      # - header
      dummy_assimilation_script_substitutions['#DUMMY_ASSIMILATION_SCRIPT_HEADER#'] = (
        models.nemo.utilities.job_script.create_script_header(machine=machine,
        partition=enkf_partition, job_name='enkf', job_time='00:02:00', exclusivity=False,
        number_of_cores=1, user_email=user_email, account=account, memory=100))
    
      # - paths
      dummy_assimilation_script_substitutions['#CLIMATETOOLBOX_PATH#'] = (
        utilities.generic.get_toolbox_root_directory())
    
      # - experiment information
      dummy_assimilation_script_substitutions['#VERSION#'] = version
      dummy_assimilation_script_substitutions['#CONFIGURATION#'] = configuration
      dummy_assimilation_script_substitutions['#EXPERIMENT_NAME#'] = name
    
      # - modules
      dummy_assimilation_script_substitutions['#CLIMATETOOLBOX_MODULE_LOADING#'] = (
        toolbox_module_loading_lines)
    
      # copy the dummy assimilation script in the experiment directory (with adjustments)
      dummy_assimilation_script_path = os.path.join(template_directory_path,
        models.nemo.utilities.job_script.DUMMY_ASSIMILATION_SCRIPT_FILE_NAME)
      updated_dummy_assimilation_script_path = os.path.join(experiment_path,
        models.nemo.utilities.job_script.DUMMY_ASSIMILATION_SCRIPT_FILE_NAME)
      dummy_assimilation_script = models.nemo.utilities.job_script.JobScript(
        file_path=dummy_assimilation_script_path)
      dummy_assimilation_script.update(substitutions=dummy_assimilation_script_substitutions,
        path=updated_dummy_assimilation_script_path)
    
    
      # write the file containing the analysis fields in the experiment directory
      analysis_fields_file_path = os.path.join(configuration_path,
        models.nemo.utilities.generic.REFERENCE_EXPERIMENT_NAME,
        models.nemo.utilities.generic.ENKF_ANALYSIS_FIELDS_STORED_FILE_NAME)
      utilities.generic.check_file_existence(path=analysis_fields_file_path,
        file_identifier='the file containing the EnKF analysis fields')
      updated_analysis_fields_file_path = os.path.join(experiment_path,
    
        models.nemo.enkf_tools.utilities.generic.ANALYSIS_FIELDS_FILE_NAME)
    
    
      with open(analysis_fields_file_path, 'r') as source, open(updated_analysis_fields_file_path,
          'w') as target:
        while True:
          line = source.readline()
    
          if line == '':
            break
    
    
          if ((models.nemo.enkf_tools.utilities.generic.OTHER_E_I_LINES_PLACEHOLDER in line) or
              (models.nemo.enkf_tools.utilities.generic.OTHER_E_S_LINES_PLACEHOLDER in line)):
            if models.nemo.enkf_tools.utilities.generic.OTHER_E_I_LINES_PLACEHOLDER in line:
              placeholder = models.nemo.enkf_tools.utilities.generic.OTHER_E_I_LINES_PLACEHOLDER
    
              number_of_layers = ice_namelist_parameters['nlay_i']
    
            else:
              placeholder = models.nemo.enkf_tools.utilities.generic.OTHER_E_S_LINES_PLACEHOLDER
    
              number_of_layers = ice_namelist_parameters['nlay_s']
    
    
            line = line.replace(placeholder, '').replace(
              models.nemo.enkf_tools.utilities.generic.NUMBER_OF_ICE_CATEGORIES_PLACEHOLDER,
    
              str(ice_namelist_parameters['jpl']))
    
    
            for layer_index in range(number_of_layers):
              target.write(line.replace('01', '%02d'%(layer_index + 1)))
    
          elif models.nemo.enkf_tools.utilities.generic.NUMBER_OF_ICE_CATEGORIES_PLACEHOLDER in line:
    
            target.write(line.replace(
              models.nemo.enkf_tools.utilities.generic.NUMBER_OF_ICE_CATEGORIES_PLACEHOLDER,
    
              str(ice_namelist_parameters['jpl'])))
    
          else:
            target.write(line)
    
      # write the sanity check namelist in the experiment directory
      sanity_check_namelist_file_path = os.path.join(template_directory_path,
        models.nemo.enkf_tools.utilities.generic.SANITY_CHECK_NAMELIST_FILE_NAME)
      updated_sanity_check_namelist_file_path = os.path.join(experiment_path,
        models.nemo.enkf_tools.utilities.generic.SANITY_CHECK_NAMELIST_FILE_NAME)
    
      with open(sanity_check_namelist_file_path, 'r') as source, open(
          updated_sanity_check_namelist_file_path, 'w') as target:
        while True:
          line = source.readline()
    
          if line == '':
            break
    
          if models.nemo.enkf_tools.utilities.generic.AMAX_N_PLACEHOLDER in line:
            target.write(line.replace(models.nemo.enkf_tools.utilities.generic.AMAX_N_PLACEHOLDER,
    
              str(ice_namelist_parameters['rn_amax_n'])))
    
          elif models.nemo.enkf_tools.utilities.generic.AMAX_S_PLACEHOLDER in line:
            target.write(line.replace(models.nemo.enkf_tools.utilities.generic.AMAX_S_PLACEHOLDER,
    
              str(ice_namelist_parameters['rn_amax_s'])))
          elif models.nemo.enkf_tools.utilities.generic.FSBC_PLACEHOLDER in line:
            target.write(line.replace(models.nemo.enkf_tools.utilities.generic.FSBC_PLACEHOLDER,
              str(ocean_namelist_parameters['nn_fsbc'])))
          elif models.nemo.enkf_tools.utilities.generic.RDT_ICE_PLACEHOLDER in line:
            target.write(line.replace(models.nemo.enkf_tools.utilities.generic.RDT_ICE_PLACEHOLDER,
    
    Antoine Barthélemy's avatar
    Antoine Barthélemy a validé
              str(ocean_namelist_parameters['rn_Dt'] * ocean_namelist_parameters['nn_fsbc'])))
    
          elif models.nemo.enkf_tools.utilities.generic.HIMIN_PLACEHOLDER in line:
            target.write(line.replace(models.nemo.enkf_tools.utilities.generic.HIMIN_PLACEHOLDER,
              str(ice_namelist_parameters['rn_himin'])))
          elif models.nemo.enkf_tools.utilities.generic.ICESAL_PLACEHOLDER in line:
            target.write(line.replace(models.nemo.enkf_tools.utilities.generic.ICESAL_PLACEHOLDER,
              str(ice_namelist_parameters['nn_icesal'])))
    
          else:
            target.write(line)
    
      # copy other files in the experiment directory
      enkf_namelist_path = os.path.join(template_directory_path,
    
        models.nemo.enkf_tools.utilities.generic.FILTER_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)
    
    
      # 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,
        restart_from_enkf=restart_from_enkf)
    
    
    
    #################
    # 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')
    
      models.nemo.enkf_tools.utilities.argument_parser.add_argument_assimilation_frequency(
        parser=parser)
    
      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)
    
      parser.add_argument(
        '-rfe', '--restart_from_enkf',
        action='store_true',
        help=('in case of restart, restart each member from the corresponding member in another EnKF '
          + 'experiment\n(default: False)'),
      )
    
      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()
    
      try:
        create_enkf_experiment(**arguments)
    
      except utilities.generic.ClimateToolboxException as e:
    
        e.display()
    
    
    if __name__ == '__main__':
      main()