source: CONFIG/publications/IPSLCM6_LMDZOR_ensembles_2023/selection/selection_GKTL_z500.py @ 6957

import xarray as xr
import numpy as np
import sys

# Paths
# create by hand the mean_weights.txt and new_trajectories.txt files
# bouger le fichier ensemble.txt juste aprÚs l'exécution du script
# créer un dossier REA dans le submit dir

# Parameters general
year = int(sys.argv[1])
name_run = sys.argv[2]
iteration = int(sys.argv[3]) - 1

path_results = "/ccc/work/cont003/gen6877/noyeller/rea_irene_rome/IPSLCM6.2.2_ENSEMBLES/config/IPSLCM6/" + name_run + "/"
path_data = "/ccc/scratch/cont003/gen6877/noyeller/IGCM_OUT/LMDZOR/DEVT/amip/" + name_run + "/ATM/Output/DA/"

# Parameters REA
k = float(sys.argv[4]) # default: -0.05
N = int(sys.argv[5]) 

# Functions
def create_name_file(name_run, year, iteration):
    date_ini_tab = ['0601','0606','0611','0616','0621','0626','0701','0706','0711','0716','0721','0726','0731','0805','0810','0815','0820','0825']
    date_end_tab = ['0605','0610','0615','0620','0625','0630','0705','0710','0715','0720','0725','0730','0804','0809','0814','0819','0824','0829']
    return name_run + "_" + str(year) + date_ini_tab[iteration-1] + "_" + str(year) + date_end_tab[iteration-1] + "_1D_histday.nc"

def compute_observable(path_data, name_file):
    data = xr.open_dataset(path_data+name_file, use_cftime=True)
    obs = data['z500'].sel(lat=slice(50,49),lon=slice(1.25,3.75)).mean(dim=('lat','lon'))
    return obs.sum(dim='time_counter').values

def compute_weights(int_obs_points, k):
    weights = np.exp(k*int_obs_points) # temporal integration in discrete time
    return weights/np.mean(weights), np.mean(weights)

def compute_cloning(weights):
    normalized_weights = weights/np.mean(weights)
    N = normalized_weights.size
    number_copies = (normalized_weights + np.random.uniform(low=0, high=1, size=N)).astype(int)
    cloning_table = np.zeros(np.sum(number_copies), dtype=int)
    indice_start = 0
    for i in range(N):
        cloning_table[indice_start:indice_start + number_copies[i]] = i
        indice_start += number_copies[i]
    delta_N = np.sum(number_copies) - N
    if delta_N >= 0 :
        return np.random.choice(cloning_table, size=N, replace=False)
    else:
        return np.random.choice(cloning_table, size=N, replace=True)

def write_ensemble_file(new_trajectories):
    with open(path_results+'ensemble.txt', 'w') as f:
        for idm,m in enumerate(new_trajectories):
            f.write("MEMBER_"+str(idm)+":"+str(m)+":1e-4")
            f.write('\n')
    f.close()

# REA

if iteration == 0:
    
    # Save mean_weights
    with open(path_results+'REA/mean_weights.txt', 'w') as f:
        f.write(str(1))
    f.close()
    
    # Save new trajectories
    new_trajectories = np.arange(N)
    with open(path_results+'REA/new_trajectories.txt', 'w') as f:
        f.write(','.join([str(n) % n for n in new_trajectories]))
    f.close()

    # Create ensemble.txt file
    write_ensemble_file(new_trajectories)

else:
    name_file = create_name_file(name_run, year, iteration)
    int_obs_points = compute_observable(path_data, name_file) # iteration in the name of name_run

    # Computations
    weights, mean_weights = compute_weights(int_obs_points, k)
    new_trajectories = compute_cloning(weights)
    print(np.max(np.unique(new_trajectories, return_counts=True)[1]))

    # Save mean_weights
    with open(path_results+'REA/mean_weights.txt', 'a') as f:
        f.write('\n')
        f.write(str(mean_weights))
    f.close()

    # Save new trajectories
    with open(path_results+'REA/new_trajectories.txt', 'a') as f:
        f.write('\n')
        f.write(','.join([str(n) % n for n in new_trajectories]))
    f.close()
    
    # Create ensemble.txt file
    write_ensemble_file(new_trajectories)