add pipeline to sf

experiments/initialize.sbatch

@@ -7,11 +7,15 @@
 experiment_id=$1
 
 echo "Initializing calibration for experiment $experiment_id"
-julia --color=no --project=experiments -e 'using Pkg; Pkg.instantiate(;verbose=true)'
+julia --color=no --project=experiments/$experiment_id -e 'using Pkg; Pkg.instantiate(;verbose=true)'
 
-julia --color=no --project=experiments -e '
+julia --color=no --project=experiments/$experiment_id -e '
     import CalibrateAtmos
     CalibrateAtmos.initialize("'$experiment_id'")
 '
 
 echo "Calibration initialized."
+
+julia --color=no --project=experiments/$experiment_id -e '
+    include("experiments/$experiment_id/model_interface.jl")
+'

experiments/model_run.sbatch

@@ -14,9 +14,10 @@ output=output/$experiment_id/$format_i/$member/model_log.out
 
 # Run the forward model
 srun --output=$output julia --color=no --project=experiments -e "
-    using ClimaComms
-    ClimaComms.init(ClimaComms.context())
-    import CalibrateAtmos
-    atmos_config = CalibrateAtmos.get_atmos_config($SLURM_ARRAY_TASK_ID, $iteration, \"$experiment_id\")
-    CalibrateAtmos.run_forward_model(atmos_config)
+    # using ClimaComms
+    # ClimaComms.init(ClimaComms.context())
+    # import CalibrateAtmos
+    include(\"experiments/$experiment_id/model_interface.jl\")
+    config = get_config($SLURM_ARRAY_TASK_ID, $iteration, \"$experiment_id\")
+    run_forward_model(config)
 "

experiments/model_run_general.sbatch

@@ -0,0 +1,22 @@
+#!/bin/bash
+#SBATCH --time=2:00:00
+#SBATCH --cpus-per-task=8
+#SBATCH --mem-per-cpu=8G
+
+# Extract command-line arguments
+experiment_id=$1
+iteration=$2
+
+# Find output directory
+format_i=$(printf "iteration_%03d" "$iteration")
+member=$(printf "member_%03d" "$SLURM_ARRAY_TASK_ID")
+output=output/$experiment_id/$format_i/$member/model_log.out
+
+# Run the forward model
+srun --output=$output julia --color=no --project=experiments -e "
+    using ClimaComms
+    ClimaComms.init(ClimaComms.context())
+    import CalibrateAtmos
+    model_config = CalibrateAtmos.get_config($SLURM_ARRAY_TASK_ID, $iteration, \"$experiment_id\")
+    CalibrateAtmos.run_forward_model(model_config)
+"

experiments/sphere_held_suarez_rhoe_equilmoist/Project.toml

@@ -0,0 +1,28 @@
+[deps]
+CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
+CalibrateAtmos = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
+ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
+ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
+ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9"
+TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
+
+[extras]
+CUDA_Driver_jll = "4ee394cb-3365-5eb0-8335-949819d2adfc"
+CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
+MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
+
+[preferences.CUDA_Driver_jll]
+compat = false
+
+[preferences.CUDA_Runtime_jll]
+local = "true"
+version = "12.2"
+
+[preferences.MPIPreferences]
+_format = "1.0"
+binary = "MPItrampoline_jll"

experiments/sphere_held_suarez_rhoe_equilmoist/generate_observations.sbatch

@@ -28,7 +28,7 @@ config = CA.AtmosConfig(config_dict)
 simulation = CA.get_simulation(config)
 CA.solve_atmos!(simulation)
 
-ta = ncread(joinpath(output_dir, "ta_60d_average.nc"), "ta")
+ta = JLD2.load(joinpath(output_dir, "ta_60d_average.nc"), "ta")
 include(joinpath(experiment_dir, "observation_map.jl"))
 (; observation, variance) = process_member_data(ta; output_variance = true)
 JLD2.save_object(joinpath(experiment_dir, "obs_mean.jld2"), observation)

src/atmos_interface.jl → experiments/sphere_held_suarez_rhoe_equilmoist/model_interface.jl

@@ -2,21 +2,21 @@ import EnsembleKalmanProcesses as EKP
 import ClimaAtmos as CA
 import YAML
 
-function get_atmos_config(member, iteration, experiment_id::AbstractString)
+function get_config(member, iteration, experiment_id::AbstractString)
     config_dict = YAML.load_file("experiments/$experiment_id/atmos_config.yml")
-    return get_atmos_config(member, iteration, config_dict)
+    return get_config(member, iteration, config_dict)
 end
 
 """
-    get_atmos_config(member, iteration, experiment_id::AbstractString)
-    get_atmos_config(member, iteration, config_dict::AbstractDict)
+    get_config(member, iteration, experiment_id::AbstractString)
+    get_config(member, iteration, config_dict::AbstractDict)
 
-Returns an AtmosConfig object for the given member and iteration. 
+Returns an AtmosConfig object for the given member and iteration.
 If given an experiment id string, it will load the config from the corresponding YAML file.
 Turns off default diagnostics and sets the TOML parameter file to the member's path.
 This assumes that the config dictionary has `output_dir` and `restart_file` keys.
 """
-function get_atmos_config(member, iteration, config_dict::AbstractDict)
+function get_config(member, iteration, config_dict::AbstractDict)
     # Specify member path for output_dir
     # Set TOML to use EKP parameter(s)
     output_dir = config_dict["output_dir"]

experiments/sphere_held_suarez_rhoe_equilmoist/observation_map.jl

@@ -27,12 +27,12 @@ function latitudinal_avg(arr)
     return dropdims(mean(arr; dims); dims)
 end
 
-function height_avg(arr)
+function height_avg(arr)  # unused
     dims = 4
     return dropdims(mean(arr; dims); dims)
 end
 
-function time_avg(arr)
+function time_avg(arr)  # unused
     dims = 1
     return dropdims(mean(arr; dims); dims)
 end

experiments/surface_fluxes_perfect_model/generate_observations.sbatch

@@ -10,8 +10,8 @@ source experiments/common_env.sh
 echo "Generating truth observations."
 
 srun  --open-mode=append julia --project=experiments -e '
-using ClimaComms
-ClimaComms.init(ClimaComms.context())
+# using ClimaComms
+# ClimaComms.init(ClimaComms.context())
 import SurfaceFluxes as SF
 import YAML
 using NetCDF
@@ -21,12 +21,12 @@ experiment_dir = joinpath("experiments", "surface_fluxes_perfect_model")
 output_dir = joinpath(experiment_dir, "truth_simulation")
 
 # file SF.surface_conditions()
-include(joinpath(experiment_dir, "calculate_fluxes.jl"))
+include(joinpath(experiment_dir, "sf_model.jl"))
 
 # observables
-ta = ncread(joinpath(output_dir, "ta_60d_average.nc"), "ta")
+ustar = ncread(joinpath(output_dir, "data/ustar_array.jld2"), "ustar_array")
 include(joinpath(experiment_dir, "observation_map.jl"))
-(; observation, variance) = process_member_data(ta; output_variance = true)
-JLD2.save_object(joinpath(experiment_dir, "obs_mean.jld2"), observation)
-JLD2.save_object(joinpath(experiment_dir, "obs_noise_cov.jld2"), variance)
+(; observation, variance) = process_member_data(ustar; output_variance = true)
+JLD2.save_object(joinpath(experiment_dir, "data/obs_mean.jld2"), observation)
+JLD2.save_object(joinpath(experiment_dir, "data/obs_noise_cov.jld2"), variance)
 '

experiments/surface_fluxes_perfect_model/generate_truth.jl

@@ -0,0 +1,53 @@
+# generate_truth: generate true y, noise and inputs
+
+FT = Float32
+include(joinpath(pathof(SF),"../../ext/CreateParametersExt.jl"))
+param_set = SurfaceFluxesParameters(FT, UF.BusingerParams) # ! this is a global variable in SF tests
+include(joinpath(pathof(SF),"../../test/test_convergence.jl"))
+
+include("sf_model.jl")
+
+
+# function to create and save "data/surface_fluxes_test_data.jld2" from generate_profiles(FT, profile_type; uf_params)
+function save_profiles(FT; profile_type =  MoistEquilProfiles(), uf_params = UF.BusingerParams, data_file = "data/surface_fluxes_test_data.jld2")
+    profiles_sfc, profiles_int = generate_profiles(FT, profile_type; uf_params)
+    data = Dict(
+        "profiles_sfc" => profiles_sfc,
+        "profiles_int" => profiles_int,
+    )
+    JLD2.save(data_file, data)
+end
+
+# function to load "data/surface_fluxes_test_data.jld2" and return profiles_sfc, profiles_int
+function load_profiles(data_file)
+    data = JLD2.load(data_file)
+    return (; profiles_sfc = data["profiles_sfc"], profiles_int = data["profiles_int"])
+end
+
+function synthetic_observed_y(x_inputs)
+    config = Dict(
+        "toml" => ["parameters.toml"],
+    )
+    y = obtain_ustar(FT, x_inputs, config)
+    # add noise to model truth to obtain y
+    Γ = 0.03^2 * I * (maximum(y) - minimum(y))
+    noise_dist = MvNormal(zeros(1), Γ)
+    apply_noise!(y, noise_dist) = y + rand(noise_dist)[1]
+    # broadcast the noise to each element of y
+    y = apply_noise!.(y, Ref(noise_dist))
+    # save y to file
+    JLD2.save("data/synthetic_observed_y.jld2", Dict("synthetic_observed_y" => y))
+    return y
+end
+
+# generate x inputs
+data_file = "data/surface_fluxes_test_data.jld2"
+save_profiles(FT, data_file = data_file)
+
+# read x inputs
+x_inputs = load_profiles(data_file)
+
+# generate synthetic observed y
+y = synthetic_observed_y(x_inputs)
+
+# inspect y

experiments/surface_fluxes_perfect_model/model_config.yml

@@ -0,0 +1,2 @@
+output_dir: output/surface_fluxes_perfect_model
+data_file: data/surface_fluxes_test_data.jld2

experiments/surface_fluxes_perfect_model/model_interface.jl

@@ -0,0 +1,60 @@
+import EnsembleKalmanProcesses as EKP
+import ClimaAtmos as CA
+import YAML
+
+include("sf_model.jl")
+
+function get_config(member, iteration, experiment_id::AbstractString)
+    config_dict = YAML.load_file("experiments/$experiment_id/atmos_config.yml")
+    return get_config(member, iteration, config_dict)
+end
+
+"""
+    get_config(member, iteration, experiment_id::AbstractString)
+    get_config(member, iteration, config_dict::AbstractDict)
+
+Returns an AtmosConfig object for the given member and iteration.
+If given an experiment id string, it will load the config from the corresponding YAML file.
+Turns off default diagnostics and sets the TOML parameter file to the member's path.
+This assumes that the config dictionary has `output_dir` and `restart_file` keys.
+"""
+function get_config(member, iteration, config_dict::AbstractDict)
+    # Specify member path for output_dir
+    # Set TOML to use EKP parameter(s)
+    output_dir = config_dict["output_dir"]
+    member_path =
+        EKP.TOMLInterface.path_to_ensemble_member(output_dir, iteration, member)
+    config_dict["output_dir"] = member_path
+    parameter_path = joinpath(member_path, "parameters.toml")
+    if haskey(config_dict, "toml")
+        push!(config_dict["toml"], parameter_path)
+    else
+        config_dict["toml"] = [parameter_path]
+    end
+
+    # Set restart file for initial equilibrium state
+    ENV["RESTART_FILE"] = config_dict["restart_file"]
+    return config_dict
+end
+
+"""
+    run_forward_model(config::AbstractDict)
+
+Runs the atmosphere model with the given an AtmosConfig object.
+Currently only has basic error handling.
+"""
+function run_forward_model(config::AbstractDict)
+
+    x_inputs = load_profiles(config.data_file)
+    obtain_ustar(FT, x_inputs, config)
+
+
+    # simulation = CA.get_simulation(config)
+    # sol_res = CA.solve_atmos!(simulation)
+    # if sol_res.ret_code == :simulation_crashed
+    #     !isnothing(sol_res.sol) && sol_res.sol .= eltype(sol_res.sol)(NaN)
+    #     error(
+    #         "The ClimaAtmos simulation has crashed. See the stack trace for details.",
+    #     )
+    # end
+end