Merge pull request #495 from SpeedyWeather/mk/summarysize

Trees with summarysize

docs/src/structure.md

@@ -107,4 +107,25 @@ field tree is the longest, defining many components for the physics parameteriza
 ```@example structure
 model = PrimitiveWetModel()
 tree(model)
+```
+
+## Size of a Simulation
+
+The `tree` function also allows for the `with_size::Bool` keyword (default `false`),
+which will also print the size of the respective branches to give you an idea of
+how much memory a SpeedyWeather simulation uses.
+
+```@example structure
+tree(simulation, max_level=1, with_size=true)
+```
+
+And with `max_level` you can truncate the tree to go down at most that many levels.
+1MB is a typical size for a one-level T31 resolution simulation. In comparison,
+a higher resolution `PrimitiveWetModel` would use
+
+```@example structure
+spectral_grid = SpectralGrid(trunc=127, nlev=8)
+model = PrimitiveWetModel(;spectral_grid)
+simulation = initialize!(model)
+tree(simulation, max_level=1, with_size=true)
 ```

src/SpeedyWeather.jl

@@ -61,17 +61,18 @@ export SpeedyTransforms, SpectralTransform
 export spectral, gridded, spectral_truncation
 include("SpeedyTransforms/SpeedyTransforms.jl")
 using .SpeedyTransforms
+import .SpeedyTransforms: prettymemory
 
 # Utility for GPU / KernelAbstractions
 include("gpu.jl")                               
 
-# abstract types
+# abstract types
 include("models/abstract_models.jl")
 include("dynamics/abstract_types.jl")
 include("output/abstract_types.jl")
 include("physics/abstract_types.jl")
 
-# GEOMETRY CONSTANTS ETC
+# GEOMETRY CONSTANTS ETC
 include("dynamics/vertical_coordinates.jl")
 include("dynamics/spectral_grid.jl")
 include("dynamics/geometry.jl")
@@ -82,7 +83,7 @@ include("dynamics/adiabatic_conversion.jl")
 include("dynamics/orography.jl")
 include("physics/land_sea_mask.jl")
 
-# VARIABLES
+# VARIABLES
 include("dynamics/particles.jl")
 include("dynamics/clock.jl")
 include("dynamics/prognostic_variables.jl")

src/models/tree.jl

@@ -5,59 +5,77 @@ $(TYPEDSIGNATURES)
 Create a tree of fields inside a Simulation instance and fields within these fields
 as long as they are defined within the modules argument (default SpeedyWeather).
 Other keyword arguments are `max_level::Integer=10`, `with_types::Bool=false`."""
-function tree(S::Simulation{M}; modules=SpeedyWeather, kwargs...) where M
-    println("Simulation{$(model_type(M))}")
-    _tree(S, modules; kwargs...)
+function tree(
+    S::Simulation{M};
+    modules=SpeedyWeather,
+    with_size::Bool = false,
+    kwargs...
+) where M
+    s = "Simulation{$(model_type(M))}"
+    s = ~with_size ? s : s*" ($(prettymemory(Base.summarysize(S))))"
+    println(s)
+    _tree(S, modules; with_size, kwargs...)
 end
 
 """
 $(TYPEDSIGNATURES)
 Create a tree of fields inside a model and fields within these fields
 as long as they are defined within the modules argument (default SpeedyWeather).
 Other keyword arguments are `max_level::Integer=10`, `with_types::Bool=false`."""
-function tree(M::ModelSetup; modules=SpeedyWeather, kwargs...)
-    println("$(model_type(M)){...}")
-    _tree(M, modules; kwargs...)
+function tree(
+    M::ModelSetup;
+    modules=SpeedyWeather,
+    with_size::Bool = false,
+    kwargs...
+)
+    s = "$(model_type(M)){...}"
+    s = ~with_size ? s : s*" ($(prettymemory(Base.summarysize(M))))"
+    println(s)
+    _tree(M, modules; with_size, kwargs...)
 end
 
 """
 $(TYPEDSIGNATURES)
 Create a tree of fields inside S and fields within these fields
 as long as they are defined within the modules argument (default SpeedyWeather).
 Other keyword arguments are `max_level::Integer=10`, `with_types::Bool=false`."""
-function tree(S; modules=SpeedyWeather, kwargs...)
-    println("$(typeof(S))")
-    _tree(S, modules; kwargs...)
+function tree(S; modules=SpeedyWeather, with_size::Bool = false, kwargs...)
+    s = "$(typeof(S))"
+    s = ~with_size ? s : s*" ($(prettymemory(Base.summarysize(S))))"
+    println(s)
+    _tree(S, modules; with_size, kwargs...)
 end
 
 function _tree(
     S,
     modules::Module...;         # fields within the modules are further inspected
-    max_level::Integer = 10,    # depth of the tree
+    max_level::Integer = 10,    # depth of the tree
     with_types::Bool = false,   # print also the types of the fields?
+    with_size::Bool = false,    # print also the Base.summarysize?
 )
-    level = 0                   # starting level of depth tree
-    prevs = falses(max_level+2) # determine whether there's still a branch levels up
-                                # needed for │ printing of upper levels that continue
+    level = 0                   # starting level of depth tree
+    prevs = falses(max_level+2) # determine whether there's still a branch levels up
+                                # needed for │ printing of upper levels that continue
     property_names = propertynames(S)
     n_properties = length(property_names)
 
     for (i,property_name) in enumerate(property_names)
         last = i == n_properties    # last elements in branches are printed with └ not ├
         prevs[level+1] = ~last
-        print_branch(property_name, S, level, max_level, last, prevs, with_types, modules...)
+        print_branch(property_name, S, level, max_level, last, prevs, with_types, with_size, modules...)
     end
 end
 
 function print_branch(
     property_name::Symbol,  # name of current field
-    S,                      # its parent struct
-    level::Integer,         # depth of tree we're on
+    S,                      # its parent struct
+    level::Integer,         # depth of tree we're on
     max_level::Integer,     # maximum depth of tree
-    last::Bool,             # is that the last field in parent?
-    prevs::BitVector,       # branching of previous fields completed?
+    last::Bool,             # is that the last field in parent?
+    prevs::BitVector,       # branching of previous fields completed?
     with_types::Bool,       # print also types
-    modules::Module...,     # in which modules to search
+    with_size::Bool,        # print object size via Base.summarysize
+    modules::Module...,     # in which modules to search
  )
     level == max_level && return nothing
 
@@ -70,7 +88,7 @@ function print_branch(
         continue_branching = true
     elseif property isa AbstractArray
         if length(property) > 0
-            property = property[1]      # unpack first element in array
+            property = property[1]      # unpack first element in array
             continue_branching = parentmodule(typeof(property)) in modules
             is_array = continue_branching
         end
@@ -84,13 +102,14 @@ function print_branch(
     junction2 = last ? "└" : "├"
     print(vertical_lines, junction2, junction1, property_name)
     s = ~with_types ? "" : "::$(typeof(property))"
+    s = ~with_size ? s : s*" ($(prettymemory(Base.summarysize(property))))"
     println(s)
 
     if continue_branching
         for (i,branch) in enumerate(property_names)
             last = i == n_properties
             prevs[level+2] = ~last
-            print_branch(branch, property, level+1, max_level, last, prevs, with_types, modules...)
+            print_branch(branch, property, level+1, max_level, last, prevs, with_types, with_size, modules...)
         end
     end
 end