Newer
Older
import os.path
Antoine Barthélemy
a validé
import utilities.generic
Antoine Barthélemy
a validé
import utilities.module_file
Antoine Barthélemy
a validé
import models.nemo.utilities.generic
Antoine Barthélemy
a validé
import models.nemo.utilities.namelist
Antoine Barthélemy
a validé
import models.nemo.utilities.job_script
Antoine Barthélemy
a validé
import models.nemo.utilities.argument_parser
import models.nemo.utilities.global_information
import models.nemo.utilities.base_experiment_plan
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
Antoine Barthélemy
a validé
import models.nemo.enkf_tools.utilities.enkf_experiment_plan
Antoine Barthélemy
a validé
#############
# variables #
#############
ASSIMILATION_SCRIPT_PYTHON_MODULE_FILES = {}
#############
# functions #
#############
def create_enkf_experiment(version, configuration, name, start_date, assimilation_start_date,
Antoine Barthélemy
a validé
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))
Antoine Barthélemy
a validé
# 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)
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,
Antoine Barthélemy
a validé
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)
Antoine Barthélemy
a validé
ice_parameters = [
Antoine Barthélemy
a validé
('jpl', int),
('nlay_i', int),
('nlay_s', int),
('rn_amax_n', float),
('rn_amax_s', float),
('rn_himin', float),
('nn_icesal', int),
Antoine Barthélemy
a validé
]
ice_namelist_parameters = ice_namelist.extract_parameters(parameters=ice_parameters)
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()
Antoine Barthélemy
a validé
# 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=' ')
Antoine Barthélemy
a validé
# 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))
Antoine Barthélemy
a validé
# - 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
Antoine Barthélemy
a validé
# - 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)
Antoine Barthélemy
a validé
Antoine Barthélemy
a validé
# copy the assimilation script in the experiment directory (with adjustments)
Antoine Barthélemy
a validé
template_directory_path = os.path.join(utilities.generic.get_toolbox_root_directory(), 'models',
'nemo', 'enkf_tools', 'templates')
Antoine Barthélemy
a validé
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)
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
# 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)
Antoine Barthélemy
a validé
# 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)
Antoine Barthélemy
a validé
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:
Antoine Barthélemy
a validé
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,
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'])))
Antoine Barthélemy
a validé
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)
Antoine Barthélemy
a validé
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)
Antoine Barthélemy
a validé
# 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,
Antoine Barthélemy
a validé
restart_frequency=restart_frequency, assimilation_frequency=assimilation_frequency)
Antoine Barthélemy
a validé
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)
Antoine Barthélemy
a validé
experiment_plan.write()
#################
# main function #
#################
def main():
print()
parser = argparse.ArgumentParser(description=create_enkf_experiment.__doc__,
formatter_class=argparse.RawTextHelpFormatter)
Antoine Barthélemy
a validé
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'
Antoine Barthélemy
a validé
+ '\n(format: %s)'%models.nemo.utilities.generic.DATE_FORMAT),
Antoine Barthélemy
a validé
models.nemo.utilities.argument_parser.add_argument_end_date(parser=parser)
models.nemo.utilities.argument_parser.add_argument_restart_frequency(parser=parser,
models.nemo.enkf_tools.utilities.argument_parser.add_argument_assimilation_frequency(
parser=parser)
Antoine Barthélemy
a validé
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'] = (
Antoine Barthélemy
a validé
models.nemo.utilities.global_information.get_version_and_configuration(
tentative_version=arguments['version'], tentative_configuration=arguments['configuration']))
Antoine Barthélemy
a validé
except utilities.generic.ClimateToolboxException as e:
e.display()
try:
create_enkf_experiment(**arguments)
Antoine Barthélemy
a validé
except utilities.generic.ClimateToolboxException as e:
e.display()
if __name__ == '__main__':
main()