run_ConFLAME.py

from train import *
from evaluate import *
from iris_plus import add_bounds

from   io     import StringIO
import numpy  as np
import cftime
import matplotlib.pyplot as plt

import datetime

import sys
sys.path.append('libs/')
sys.path.append('src/attribution/')
from climtatology_difference import *
from above_percentile_mean import *
from plot_maps import *
from plot_multimaps import *
from attribution import *

try:
    from concurrent.futures import ProcessPoolExecutor, as_completed
    from multiprocessing import get_context
except:
    pass

def call_eval(training_namelist, namelist,
              control_run_name, extra_params = None, run_only = True, *args, **kw):
    
    return evaluate_MaxEnt_model_from_namelist(training_namelist, namelist,
                                               run_only = run_only, 
                                               control_run_name = control_run_name,
                                               extra_params = extra_params, *args, **kw)
    
   
def Standard_limitation(training_namelist, namelist,
                        controlID, name, control_direction, *args, **kws):  
    
    control_Directioni = np.array(control_direction.copy())
    control_Directioni[:] = 0.0
    control_Directioni[controlID] = control_direction[controlID]
    
    extra_params = {"control_Direction": control_Directioni}
    
    return call_eval(training_namelist, namelist,
                     name + '/Standard_'+ str(controlID), extra_params, hyper = False,
                     *args, **kws)

def Potential_limitation(training_namelist, namelist,
                        controlID, name, control_direction, *args, **kws):   
    control_Directioni = np.array(control_direction.copy())
    control_Directioni[controlID] = 0.0
    
    extra_params = {"control_Direction": control_Directioni}
    
    return call_eval(training_namelist, namelist,
                     name + '/Potential'+ str(controlID), extra_params, hyper = False,
                     *args, **kws)



def Without_limitation(training_namelist, namelist,
                        controlID, name, control_direction, *args, **kws):   
    control_Directioni = np.array(control_direction.copy())
    control_Directioni[controlID] = 0.0
    
    extra_params = {"control_Direction": control_Directioni}
    
    return call_eval(training_namelist, namelist,
                     name + '/Without'+ str(controlID), extra_params, hyper = False,
                     *args, **kws)

def Potential_limitation(training_namelist, namelist,
                        controlID, name, control_direction, *args, **kws):   
    Others = Without_limitation(training_namelist, namelist,
                              controlID, name, control_direction,
                              *args, **kws)[0]   
    Control = call_eval(training_namelist, namelist,
                     name + '/control', hyper = False,
                     *args, **kws)

    def for_realization(r):
        ens_no = Control.coord('realization').points[r]
        filename = info['dir_outputs'] + '/samples/' + \
                   info['filename_out'] + '/' + name + \
                   '/Potential_' + str(controlID) + '/'
        makeDir(filename) 
        filename = filename  + 'sample-pred' + str(ens_no) + '.nc'
        if os.path.exists(filename):
            return iris.load_cube(filename)
        
        out = Other[r] - Control[r]
        iris.save(out, filename)
        return out
    
    outs = [for_realization(r) for r in range(Control.shape[0])]
    outs = iris.cube.CubeList(outs).merge_cube()
    
    return  [outs]

def Potential_climateology_limitation(training_namelist, namelist,
                        controlID, name, control_direction, *args, **kws):   
    
    info = read_variables_from_namelist(training_namelist) 
    Control = Standard_limitation(training_namelist, namelist,
                                  controlID, name, control_direction,
                                  *args, **kws)[0]   

    control_Directioni = np.array(control_direction.copy())
    control_Directioni[controlID] = 0.0
    
    extra_params = {"control_Direction": control_Directioni}
    
    Others = Without_limitation(training_namelist, namelist, controlID, name, control_direction,*args, **kws)[0]   

    def for_realization(r):
        ens_no = Control.coord('realization').points[r]
        filename = info['dir_outputs'] + '/samples/' + \
                   info['filename_out'] + '/' + name + \
                   '/Potential_climatology' + str(controlID) + '/'
        makeDir(filename) 
        filename = filename  + 'sample-pred' + str(ens_no) + '.nc'
        if os.path.exists(filename):
            return iris.load_cube(filename)

        anomaly_cube = Control[r].copy()
        Climateology = anomaly_cube.aggregated_by('month', iris.analysis.MEAN)
        for month in range(1, 13):
            # Make a mask for the current month
            month_mask = anomaly_cube.coord('month').points == month
            
            # Get climatology slice for this month
            clim_slice = Climateology.extract(iris.Constraint(month=month))
            
            # Subtract climatology from all matching time steps
            anomaly_cube.data[month_mask, :] -= clim_slice.data[None, :, :]
     
        out = anomaly_cube * Others[r]
        iris.save(out, filename)
        return out
    
    outs = [for_realization(r) for r in range(Control.shape[0])]
    outs = iris.cube.CubeList(outs).merge_cube()
    
    return [outs]

  
def make_time_series(cube, name, output_path, percentile = None, cube_assess = None, 
                     grab_old = False, *args, **kw):
    
    print("finding " + str(percentile) + " for " + name + "\n\t into:" + output_path)
    print(datetime.datetime.now())
    if percentile is None or percentile == 0.0:        
        out_dir = output_path + '/mean/'
    else:
        out_dir = output_path + '/pc-' + str(percentile) + '/'
    
    lock_file = out_dir + 'lock/' 
    makeDir(lock_file)
    lock_file += name + '.txt'
     
    if os.path.isfile(lock_file) and grab_old:    
        return out_dir
    
    if cube_assess is None: cube_assess = cube
    cube = add_bounds(cube)
    cube_assess = add_bounds(cube_assess)
    
    cube.data = np.ma.masked_invalid(cube.data)
    grid_areas = iris.analysis.cartography.area_weights(cube)
       
    if percentile is None or percentile == 0.0:
        area_weighted_mean =  [cube[i].collapsed(['latitude', 'longitude'],
                                                 iris.analysis.MEAN, weights = grid_areas[i]) \
                                   for i in range(cube.shape[0])]
        
        area_weighted_mean = iris.cube.CubeList(area_weighted_mean).merge_cube().data
        if area_weighted_mean.ndim == 1:
            area_weighted_mean = area_weighted_mean.reshape([1, len(area_weighted_mean)])
    
    else:
        def percentile_for_relization(cube, i = None):
            print("\tprocessing enemble" + str(i))
            if i is not None:
                cubei = cube[i]
                cube_assessi = cube_assess[i]
                ntime =  range(cube.shape[1])
            else:
                cubei = cube
                cube_assessi = cube_assess
                ntime =  range(cube.shape[0])
                
            out = [above_percentile_mean(cubei[j], cube_assessi[j], percentile, *args, **kw) \
                   for j in ntime]
            return out
        
        if cube.ndim == 4:
            area_weighted_mean = np.array([percentile_for_relization(cube, i) \
                                          for i in range(cube.shape[0])])  
        else:
            
            area_weighted_mean =np.array([percentile_for_relization(cube, None)])
    
    
    climatology, anomaly, ratio = climtatology_difference(area_weighted_mean)
    makeDir(out_dir)
    
    def output_cube_to_csv(data, realizations, extra_dim, filename): 
        
        times = cube.coord('time').units.num2date(cube.coord('time').points)[0:data.shape[1]]
        try:
            df = pd.DataFrame(data, index=realizations, columns=[t.isoformat() for t in times])
        except:
            df = pd.DataFrame(data[:,0:len(times)], 
                              index=realizations, columns=[t.isoformat() for t in times])
            
        df.index.name = extra_dim
        df.to_csv(filename)
        #np.savetxt(out_file_points, area_weighted_mean.data, delimiter=',')

    percentiles = [5, 10, 25, 50, 75, 90, 95]

    def make_output_TS(data, dir = ''):
        
        out_file_points = out_dir + '/members/' + dir + '/'
        out_file_TS = out_dir + '/percentles/' + dir + '/' 
        makeDir(out_file_points)
        makeDir(out_file_TS)
        
        rownames = cube.coord('realization').points
        if len(rownames) == 1: 
            rownames = [0]
        
        output_cube_to_csv(data, rownames, 
                       'realization', out_file_points  + name + '.csv')
        TS = np.nanpercentile(data, percentiles, axis = 0)
        output_cube_to_csv(TS, percentiles,  'percentiles', out_file_TS  + name + '.csv')
    
    
    make_output_TS(area_weighted_mean, 'absolute')
    make_output_TS(climatology, 'climatology')
    make_output_TS(anomaly, 'anomaly')
    make_output_TS(ratio, 'ratio')
    
    os.makedirs(os.path.dirname(lock_file), exist_ok=True)
    Path(lock_file).touch()
    return out_dir

def make_both_time_series(percentiles, *args, **kw):
    if percentiles is None: return None
    for percentile in percentiles:
        make_time_series(*args, **kw, percentile = percentile) 


def run_experiment(training_namelist, namelist, control_direction,
                   control_names, control_colours, plot_control_maps,
                   output_dir, output_file, fig_dir, 
                   name = '', time_series_percentiles = None, 
                   limitation_types = None, controls_to_plot = None,*args, **kws):
    
    if "baseline" in name: 
        run_only = False
    else:
        run_only = True
    print("running: " + name)
    name = name + '-'
    
    temp_file = 'temp/run_ConFire_lock' + (output_dir + \
            output_file + name).replace('/', '_') + '.txt'
    #if os.path.isfile(temp_file): return None

    figName = fig_dir + '/' + name + 'control_TS'
    makeDir(figName + '/')
    
    Evaluate, Y, X, lmask, scalers  = call_eval(training_namelist, namelist,
                        name + '/Evaluate', run_only = run_only, return_inputs = True,
                        filename_out_ext = 'stochastic', fig_dir = fig_dir, 
                        *args, **kws)

    Control, Y, X, lmask, scalers  = call_eval(training_namelist, namelist,
                        name + '/control', run_only = run_only, return_inputs = True, 
                        Y = Y, X = X, lmask = lmask, scalers = scalers,
                        sample_error = False, filename_out_ext = 'none_stochastic',
                        fig_dir = fig_dir,
                        *args, **kws)
    
        
    grab_old = read_variables_from_namelist(namelist)['grab_old_trace']
    out_dir_ts = output_dir +'/time_series/' +  output_file + '/' + name

    try:
        obs = Control[1].copy()
        obs.data[~obs.data.mask] = Y
        out_dir_samples = output_dir + '/samples/' +  output_file + '/' + name +  \
                        '/observation.nc'
    
        iris.save(obs, out_dir_samples)
        make_both_time_series(time_series_percentiles, obs, 'observation', 
                              out_dir_ts, grab_old = grab_old)
    except:
        pass
    
    evaluate_TS = make_both_time_series(time_series_percentiles, Evaluate[0], 'Evaluate', 
                                        out_dir_ts,
                                        cube_assess = Control[0], grab_old = grab_old)

    
    
    control_TS = make_both_time_series(time_series_percentiles, Control[0], 'Control', 
                                       out_dir_ts,
                                       cube_assess = Control[0], grab_old = grab_old)
    
    if limitation_types is not None and controls_to_plot is not None:
        if control_names is None: 
            control_names = [srt(i) for i in controls_to_plot]
        limitation_types_funs = []
        for ltype in limitation_types:
            if ltype == 'standard':
                limitation_types_funs += [Standard_limitation]
            if ltype == 'potential':
                limitation_types_funs += [Potential_limitation]
            if ltype == 'potential_climateology':
                limitation_types_funs += [Potential_climateology_limitation]
         
        for ltype, FUN in zip(limitation_types,limitation_types_funs):
            limitation = [FUN(training_namelist, namelist, i, 
                          name, control_direction, *args, 
                          Y = Y, X = X, lmask = lmask, scalers = scalers, 
                              cube_assess = Control[0], **kws) \
                        for i in controls_to_plot]
            
            if plot_control_maps:
                plot_ensemble_maps(limitation, titles = control_names,
                                   control_colours = control_colours,
                                   output_path = fig_dir + \
                                            ltype + 'controls_maps.png')

            limitation_TS = np.array([make_both_time_series(time_series_percentiles, \
                                                        cube[0], \
                                                        ltype + '-' + name, out_dir_ts, \
                                                        grab_old = grab_old) \
                               for cube, name in zip(limitation, control_names)])
        
    open(temp_file, 'a').close() 

    

def run_experiment_wrapper(kwargs):
    run_experiment(**kwargs)
    try:
        run_experiment(**kwargs)
        return (kwargs, "success")
    except Exception as e:
        print(f"[ERROR] Experiment failed with kwargs={kwargs['name']}:\n{e}")
        import traceback
        traceback.print_exc()
        return (kwargs, f"error: {e}")

def run_ConFire(namelist):   
    print(f"Running ConFLAME with namelist: {namelist}")
    run_info = read_variables_from_namelist(namelist) 

    def select_from_info(item, alternative = None):
        try:
            out = run_info[item]
        except:
            out = alternative
        return out

    control_direction = select_from_info('control_Direction')
    
    if control_direction is None:
        control_direction = [param['value'] for param in run_info['priors'] \
                             if param['pname'] == 'control_Direction'][-1]
    
    control_names = select_from_info('control_names')
    control_colours = select_from_info('control_colours')
    plot_control_maps = select_from_info('plot_control_maps', False)
    subset_function_args = select_from_info('subset_function_args')
    subset_function_eval = select_from_info('subset_function_eval')
    subset_function_args_eval = select_from_info('subset_function_args_eval')
    if subset_function_args_eval is None: subset_function_args_eval =subset_function_args 
    regions = select_from_info('regions')
    time_series_percentiles = select_from_info('time_series_percentiles')
     
    def run_for_regions(region = None):
        
        if region is None:
            region = '<<region>>'
        else:
            def set_region_months(ssa):
                if isinstance(ssa, list):
                    for i in range(len(ssa)):
                        try:
                            ssa['months_of_year'] = run_info['region_months'][region]
                        except:
                            pass
                else:
                    ssa['months_of_year'] = run_info['region_months'][region]
                return ssa
            if select_from_info('region_mnths') is not None:
                set_region_months(subset_function_args)
                set_region_months(subset_function_args_eval)
        model_title = run_info['model_title'].replace('<<region>>', region)
        dir_training = run_info['dir_training'].replace('<<region>>', region)
        if 'dir_projecting' in run_info.keys():
            dir_projecting = run_info['dir_projecting'].replace('<<region>>', region)
        else:
            dir_projecting = dir_training
        trace, scalers, training_namelist = \
                        train_MaxEnt_model_from_namelist(namelist, model_title = model_title,
                                                         dir_training = dir_training,
                                                         subset_function_args = subset_function_args)
        params = read_variables_from_namelist(training_namelist)
        output_dir = params['dir_outputs']
        output_file = params['filename_out']
        fig_dir = output_dir + '/figs/' + output_file + '/'
        os.makedirs(fig_dir, exist_ok=True)
        def find_replace_period_model(exp_list):
            exp_list_all = [item.replace('<<region>>', region) for item in exp_list \
                            if "<<experiment>>" not in item and "<<model>>" not in item]
            looped_items = [item for item in exp_list \
                            if "<<experiment>>" in item and "<<model>>" in item]
            if periods is None and periods is None: 
                return exp_list_all
            for experiment, period in zip(experiments, periods):
                for model in models:
                    dirs = [item.replace("<<period>>", period) for item in looped_items]    
                    dirs = [item.replace("<<model>>", model) for item in dirs]   
                    dirs = [item.replace("<<experiment>>", experiment) for item in dirs] 
                    dirs = [item.replace('<<region>>', region) for item in dirs] 
                    exp_list_all += dirs
             
            return exp_list_all
        
        y_filen = [run_info['y_filen']]
        names_all = ['baseline']
        exp_type = ['single']        
        dirs_all = [params['dir']]
        common_noises = [True]
        limitation_types = select_from_info('limitation_types')
        max_no_ensembles =  select_from_info('max_no_ensembles')
        try:
            y_filen1 = [select_from_info('y_filen_eval', run_info['x_filen_list'][0])]
            experiment_dirs  = select_from_info('experiment_dir')
            experiment_names = select_from_info('experiment_names')
            experiments = select_from_info('experiment_experiment')
            periods = select_from_info('experiment_period')
            models = select_from_info('experiment_model')
            controls_to_plot = select_from_info('controls_to_plot', 
                                                 range(len(control_direction)))
            experiment_dirs = find_replace_period_model(experiment_dirs)
            experiment_names = find_replace_period_model(experiment_names)
            exp_type = exp_type + \
                select_from_info('experiment_type', ['single'] * len(experiment_names))
            names_all = names_all + experiment_names
            dirs_all = dirs_all + experiment_dirs
            y_filen = y_filen + y_filen1 * len(experiment_dirs)
            common_noises = common_noises + \
                select_from_info('experiment_common_noise',[True] * len(experiment_names))
            
        except:
            pass   
        
        args_list = [dict(training_namelist=training_namelist,
                          namelist=namelist,
                          control_direction=control_direction,
                          control_names=control_names,
                          control_colours=control_colours,
                          plot_control_maps = plot_control_maps,
                          output_dir=output_dir,
                          output_file=output_file,
                          fig_dir = fig_dir,
                          name=name,
                          time_series_percentiles=time_series_percentiles,
                          limitation_types = limitation_types, 
                          controls_to_plot = controls_to_plot,
                          dir=dir,
                          experiment_type = expt,
                          y_filen=yfile,
                          model_title=model_title,
                          subset_function = subset_function_eval,
                          subset_function_args = subset_function_args_eval,
                          max_no_ensembles = max_no_ensembles,
                          common_noise = common_noise
                         )
                    for name, dir, expt, yfile, common_noise \
                        in zip(names_all, dirs_all, exp_type, y_filen, common_noises)
                ]
        #args_list.reverse()

        
        if len(args_list) > 1 and select_from_info('parallelize', True): 
            try:
                with get_context("spawn").Pool(processes=4) as pool:
                    pool.map(run_experiment_wrapper, args_list)
            except:
                for args in args_list:
                    run_experiment_wrapper(args)
        else:
            for args in args_list:
                run_experiment_wrapper(args)

        #if len(args_list)>1:
        #    #try:    
        #    attribution_analysis(output_dir, '/' + output_file + '/', 
        #                             [x["dir"] for x in args_list if x["name"] == "factual"][0],
        #                             obs_file_nc = args_list[0]['y_filen'], out_dir = fig_dir)
        #    #except:
        #    #    pass
    if regions is None:
        run_for_regions(None)
    else:
        for region in regions: run_for_regions(region)

if __name__=="__main__":

    if len(sys.argv) < 2:
        print("Usage: python run_ConFire.py <namelist_path>")
        sys.exit(1)
    namelist = sys.argv[1]

    run_ConFire(namelist)