can't initialize a column model with multiple longitude points

[brian-rose]

There is currently no num_lon option for the climlab.column_state() convenience function for setting up grids for atmospheric models.

E.g. this should work but it doesn't:

import climlab
state = climlab.column_state(num_lon=10)

[HenryDane]

I've added a volume_state function to make construction of 3D volumes easier. I can submit a pull request for this feature alone if you would like.

As far as I can tell, there are some challenges to getting 3D states working with climlab, mainly that RRTMG does not work reliably in the presence of latitude, longitude, and altitude at once. I've attempted some changes to RRTMG to fix the problems but it produces non-physical results.

I've detailed some of the investigation below in case it is helpful down the line. Most of the dicussion closely relates to #50 .

---

One of the major obstacles to implementing 3D states is that Field objects have an axis_index object which allows them to track which axis is associated with which numpy axis. This, on its own, is good, but climlab does not generally respect this and tends to use hard-coded axis indices. This causes some complications since the order of axes is not guaranteed (although climlab does generally adhere to a convention for axis order).

Other techical challenges come in similarly subtle forms. In the _prepare_general_arguments function, there are some lines like:

o3vmr    = _climlab_to_rrtm(RRTMGobject.absorber_vmr['O3'] * np.ones_like(RRTMGobject.Tatm))

which seem on the surface to be fine but do not work correctly in the presence of all three axes. A potential fix for this (which is not very elegant) is:

def _safe_avmr_mult(a, b):
    if np.isscalar(a) or np.isscalar(b):
        return a * b
    if len(a.shape) == len(b.shape):
        return a * b
    '''Multiplies lat,lev by lat,lon,lev'''
    return np.repeat((a * np.ones_like(b[:,0,:]))[:,np.newaxis,:], b.shape[1], axis=1)

# ...

o3vmr    = _climlab_to_rrtm(_safe_avmr_mult(RRTMGobject.absorber_vmr['O3'], np.ones_like(RRTMGobject.Tatm)))

A lot of the subtle issues with array index handling could be solved by overriding the __mul__ (and probably also __rmul__) function of the Field class and making that overridden version guarantee that values are multiplied together along correct axes. This approach could also be used for other operators.

Overall, I think that the approach of migrating more heavily to using Field (or something similar) would really help.

[brian-rose]

@HenryDane I have some work in progress (though honestly not really progressing much) that completely refactors the climlab internals to use xarray objects for all model fields and get rid of the Field class entirely.

You're absolutely correct that climlab in its current form is inconsistent in its usage of the Field class. That's the messy legacy of me trying to invent a system for labelled gridded data objects prior to the existence of Xarray.

By switching the internals to Xarray data array objects, we will lose some performance due to overhead, but outsource a whole lot of array dimension logic and drastically simplify the climlab code. I think it's long overdue but I haven't found the time to pursue it as far as it needs.