<commit_before># coding: utf-8
"""
The payu interface for the UM atmosphere model
-------------------------------------------------------------------------------
Contact: <NAME> <<EMAIL>>
-------------------------------------------------------------------------------
Distributed as part of Payu, Copyright 2011 <NAME>
Licensed under the Apache License, Version 2.0
http://www.apache.org/licenses/LICENSE-2.0
"""

from payu.modeldriver import Model


class UnifiedModel(Model):

    #---
    def __init__(self, expt, name, config):
        super(UnifiedModel, self).__init__(expt, name, config)

        self.model_type = 'um'
        self.default_exec = 'um'

        self.modules = ['pbs',
                        'openmpi']


    #---
    def archive(self):
        raise NotImplementedError

    def set_model_pathnames(self):
        super(UnifiedModel, self).set_model_pathnames()

        ice_nml_path = os.path.join(self.control_path, self.ice_nml_fname)
        self.ice_nmls = f90nml.read(ice_nml_path)

        res_path = os.path.normpath(setup_nml['restart_dir'])
        if not os.path.isabs(res_path):
            res_path = os.path.join(self.work_path, res_path)

        self.work_init_path = res_path
        self.work_restart_path = res_path

        work_out_path = os.path.normpath(setup_nml['history_dir'])
        if not os.path.isabs(work_out_path):
            work_out_path = os.path.join(self.work_path, work_out_path)

        self.work_output_path = work_out_path

        # Determine if there is a work input path
        grid_nml = self.ice_nmls['grid_nml']
        input_path, grid_fname = os.path.split(grid_nml['grid_file'])
        if input_path and not input_path == '.':
            assert not os.path.isabs(input_path)


        self.work_input_path = os.path.join(self.work_path, input_path)


    #---
    def set_model_output_paths(self):
        super(UnifiedModel, self).set_model_output_paths()

        res_dir = self.ice_nmls['setup_nml']['restart_dir']

        # Use the local initialization restarts if present
        # TODO: Check for multiple res_paths across input paths?
        if self.expt.counter == 0:
            for input_path in self.input_paths:
                if os.path.isabs(res_dir):
                    init_res_path = res_dir
                else:
                    init_res_path = os.path.join(input_path, res_dir)
                if os.path.isdir(init_res_path):
                    self.prior_restart_path = init_res_path


    #---
    def get_prior_restart_files(self):
        return [f for f in os.listdir(self.prior_restart_path)
                if f.startswith('iced.')]

    #---
    def collate(self):
        raise NotImplementedError

    #---
    def setup(self):
        super(UnifiedModel, self).setup()

        if self.expt.counter > 0:

            # Update input namelist
            setup_nml = self.ice_nmls['setup_nml']

            setup_nml['runtype'] = 'continue'
            setup_nml['restart'] = True

            nml_path = os.path.join(self.work_path, self.ice_nml_fname)
            f90nml.write(self.ice_nmls, nml_path + '~')
            shutil.move(nml_path + '~', nml_path)

            # Generate ice.restart_file
            # TODO: Check the filenames more aggressively
            last_restart_file = sorted(self.get_prior_restart_files())[-1]

            res_ptr_path = os.path.join(self.work_init_path, 'ice.restart_file')
            with open(res_ptr_path, 'w') as res_ptr:
                print(last_restart_file, file=res_ptr)


    #---
    def archive(self, **kwargs):

        for f in os.listdir(self.work_input_path):
            f_path = os.path.join(self.work_input_path, f)
            if os.path.islink(f_path):
                os.remove(f_path)

        # Archive restart files before processing model output
        cmd = 'mv {src} {dst}'.format(src=self.work_restart_path,
                                      dst=self.restart_path)
        rc = sp.Popen(shlex.split(cmd)).wait()
        assert rc == 0
<commit_msg>Set up environment variables for the UM.<commit_after># coding: utf-8
"""
The payu interface for the UM atmosphere model
-------------------------------------------------------------------------------
Contact: <NAME> <<EMAIL>>
-------------------------------------------------------------------------------
Distributed as part of Payu, Copyright 2011 <NAME>
Licensed under the Apache License, Version 2.0
http://www.apache.org/licenses/LICENSE-2.0
"""

from payu.modeldriver import Model
import os
import sys
import imp
import shlex
import shutil

class UnifiedModel(Model):

    #---
    def __init__(self, expt, name, config):
        super(UnifiedModel, self).__init__(expt, name, config)

        self.model_type = 'um'
        self.default_exec = 'um'

        self.modules = ['pbs',
                        'openmpi']

    #---
    def archive(self):
        raise NotImplementedError

    def set_model_pathnames(self):
        super(UnifiedModel, self).set_model_pathnames()

        ice_nml_path = os.path.join(self.control_path, self.ice_nml_fname)
        self.ice_nmls = f90nml.read(ice_nml_path)

        res_path = os.path.normpath(setup_nml['restart_dir'])
        if not os.path.isabs(res_path):
            res_path = os.path.join(self.work_path, res_path)

        self.work_init_path = res_path
        self.work_restart_path = res_path

        work_out_path = os.path.normpath(setup_nml['history_dir'])
        if not os.path.isabs(work_out_path):
            work_out_path = os.path.join(self.work_path, work_out_path)

        self.work_output_path = work_out_path

        # Determine if there is a work input path
        grid_nml = self.ice_nmls['grid_nml']
        input_path, grid_fname = os.path.split(grid_nml['grid_file'])
        if input_path and not input_path == '.':
            assert not os.path.isabs(input_path)


        self.work_input_path = os.path.join(self.work_path, input_path)


    #---
    def set_model_output_paths(self):
        super(UnifiedModel, self).set_model_output_paths()

        res_dir = self.ice_nmls['setup_nml']['restart_dir']

        # Use the local initialization restarts if present
        # TODO: Check for multiple res_paths across input paths?
        if self.expt.counter == 0:
            for input_path in self.input_paths:
                if os.path.isabs(res_dir):
                    init_res_path = res_dir
                else:
                    init_res_path = os.path.join(input_path, res_dir)
                if os.path.isdir(init_res_path):
                    self.prior_restart_path = init_res_path


    #---
    def get_prior_restart_files(self):
        return [f for f in os.listdir(self.prior_restart_path)
                if f.startswith('iced.')]

    #---
    def collate(self):
        raise NotImplementedError

    #---
    def setup(self):
        super(UnifiedModel, self).setup()

        # Set up environment variables needed to run UM. 
        # Look for a python file in the config directory.
        um_env = imp.load_source('um_env',
                os.path.join(self.control_path, 'um_env.py'))

        # Put all in the current environment. 
        os.environ.update(um_env.vars)
{'fraction_non_alphanumeric': 0.14327162199502624, 'fraction_numerical': 0.004006631666206135, 'mean_word_length': 1.9777869189633897, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '465119', 'n_tokens_mistral': 2332, 'n_tokens_neox': 2215, 'n_words': 451}