save function crashing with multiple lompe objects of the same type

save_model function error occurrence when model inputs include multiple data objects of the same type. In this instance I was trying to use two separate convection inputs (one for a SuperDARN data set, one for a PFISR data set) from the same experiment/event timeframe. The below code in the MWE works for just PFISR or just SuperDARN as inputs, but produces the below error when trying to use both simultaneously.

Error:

 Traceback (most recent call last):

  File ~/miniconda3/envs/lompe/lib/python3.11/site-packages/spyder_kernels/py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File ~/Projects/isr_3d_vefs/pfrr_runscript.py:29
    run_lompe_pfisr(start_time, end_time, time_step, Kp, x_resolution,

  File ~/Projects/isr_3d_vefs/TEST_pfrr_run_lompe.py:153 in run_lompe_pfisr
    save_model(model, file_name=savefile) # one file per time stamp

  File ~/miniconda3/envs/lompe/lib/python3.11/site-packages/lompe/utils/save_load_utils.py:181 in save_model
    if data_locs: data_vars1.update(data_locs_to_dict(model))

  File ~/miniconda3/envs/lompe/lib/python3.11/site-packages/lompe/utils/save_load_utils.py:351 in data_locs_to_dict
    coords= np.array(coords)

ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (2, 2) + inhomogeneous part.

MWE (pseudo code - functions that collect data are not included):

    # set up grid
    position = (-147, 65) # lon, lat
    orientation = (-1, 2) # east, north
    L, W, Lres, Wres = 500e3, 500e3, x_resolution, y_resolution # dimensions and resolution of grid
    grid = lompe.cs.CSgrid(lompe.cs.CSprojection(position, orientation), L, W, Lres, Wres, R = 6481.2e3)

    # set up conductances and model
    SH = lambda lon = grid.lon, lat = grid.lat: hardy_EUV(lon, lat, Kp, time_intervals[0], 'hall'    )
    SP = lambda lon = grid.lon, lat = grid.lat: hardy_EUV(lon, lat, Kp, time_intervals[0], 'pedersen')
    model = lompe.Emodel(grid, Hall_Pedersen_conductance = (SH, SP))

    # Collect datasets here (uses outside functions to collect data)
    pfisr_data = pfisr.collect_data(pfisrfn, time_intervals)
    mag_data = mag.collect_data(pokermagfn, time_intervals)
    superdarn_ksr_data = sd.collect_data(superdarn_direc, time_intervals, 'ksr/')

    for i, (stime, etime) in enumerate(time_intervals):
        t = stime
        print("t: ",t)

        SH = lambda lon = grid.lon, lat = grid.lat: hardy_EUV(lon, lat, Kp, t, 'hall'    )
        SP = lambda lon = grid.lon, lat = grid.lat: hardy_EUV(lon, lat, Kp, t, 'pedersen')

        model.clear_model(Hall_Pedersen_conductance = (SH, SP)) # reset

        # add datasets for this time
        model.add_data(pfisr_data[i])
        model.add_data(mag_data[i])
        model.add_data(superdarn_ksr_data[i])

        # run model
        gtg, ltl = model.run_inversion(l1 = 2, l2 = 0.1)

        # USE FOR SAVING MODEL NCs
        savefile = "test_output" # create directory to save output as nc to read in

        save_model(model, file_name=savefile) # one file per time stamp

klaundal / lompe

save function crashing with multiple lompe objects of the same type #39