CliMA / ClimaAtmos.jl

ClimaAtmos.jl is a library for building atmospheric circulation models that is designed from the outset to leverage data assimilation and machine learning tools. We welcome contributions!
Apache License 2.0
82 stars 16 forks source link

Make AtmosModel more structured #3385

Open charleskawczynski opened 1 week ago

charleskawczynski commented 1 week ago

I looked at specializations (following https://timholy.github.io/SnoopCompile.jl/dev/tutorials/pgdsgui/), and it seems that julia is overspecializing the AtmosModel, likely because there are too many type parameters:

julia> t, m = tms[end-7]
(5.9745849170000005, var"#AtmosModel#36"(model_config, moisture_model, precip_model, cloud_model, call_cloud_diagnostics_per_stage, forcing_type, subsidence, ozone, radiation_mode, ls_adv, external_forcing, edmf_coriolis, advection_test, tendency_model, edmfx_model, turbconv_model, non_orographic_gravity_wave, orographic_gravity_wave, hyperdiff, vert_diff, diff_mode, sgs_adv_mode, viscous_sponge, rayleigh_sponge, sfc_temperature, insolation, surface_model, surface_albedo, numerics, ::Type{ClimaAtmos.AtmosModel}) @ ClimaAtmos ~/Dropbox/Caltech/work/dev/CliMA/ClimaAtmos.jl/src/solver/types.jl:392)

As a result, creating this object takes 5.97 seconds, which seems pretty long.

We should probably convert if from a flat datastructure to a more structured one. For example, what I'd preliminarily propose is:

    config
        numerics::NUM = nothing
        model_config::MC = nothing
    moisture::MM = nothing
        moisture_model::MM = nothing
        precip_model::PM = nothing
        cloud_model::CM = nothing
            call_cloud_diagnostics_per_stage::CCDPS = nothing
    forcing::F = nothing
        forcing_type::F = nothing
        subsidence::S = nothing
        external_forcing::EXTFORCING = nothing
    radiation::RM = nothing
        radiation_mode::RM = nothing
        ozone::OZ = nothing
        insolation::IN = nothing
    advection::A
        ls_adv::LA = nothing
        advection_test::AT = nothing
    tendency_model::TM = nothing
    turbconv::TCM = nothing
        edmf_coriolis::EC = nothing
        edmfx_model::EDMFX = nothing
        turbconv_model::TCM = nothing
        sgs_adv_mode::SAM = nothing
    gravity_wave::NOGW = nothing
        non_orographic_gravity_wave::NOGW = nothing
        orographic_gravity_wave::OGW = nothing
    hyperdiff::HD = nothing
    vert_diff::VD = nothing
        diff_mode::DM = nothing
    sponge::SP = nothing
        viscous_sponge::VS = nothing
        rayleigh_sponge::RS = nothing
    surface::SU
        sfc_temperature::ST = nothing
        surface_model::SM = nothing
        surface_albedo::SA = nothing

This would reduce type parameters of AtmosModel from 29 to 12.

Sbozzolo commented 1 week ago

See also https://github.com/CliMA/ClimaAtmos.jl/issues/2063