Bug in Elementary charges in json loading

[CalCraven]

Seems to be an issue in saving and loading a json GMSO object with charges that are elementary charges.

Code to reproduce issue:

import gmso
top = gmso.Topology.load(f"top.lammps")
print(top.sites[0].charge
top.save("top.json", overwrite=True)
top = gmso.Topology.load(f"top.json")

Output

-0.834 elementary_charge
---------------------------------------------------------------------------
UnitParseError                            Traceback (most recent call last)
Cell In[71], line 5
      3 print(top.sites[0].charge)
      4 top.save("source_files/top.json", overwrite=True)
----> 5 top = gmso.Topology.load(f"source_files/top.json")
      6 for fn in ["gro", "lammps", "json", "xyz", "mcf"]:
      7     print(fn)

File ~/Dropbox/Mac/Documents/Vanderbilt/Research/MoSDeF/gmso/gmso/core/topology.py:1825, in Topology.load(cls, filename, **kwargs)
   1822 from gmso.formats import LoadersRegistry
   1824 loader = LoadersRegistry.get_callable(filename.suffix)
-> 1825 return loader(filename, **kwargs)

File ~/Dropbox/Mac/Documents/Vanderbilt/Research/MoSDeF/gmso/gmso/formats/json.py:326, in load_json(filename)
    324 with filename.open("r") as json_file:
    325     json_dict = json.load(json_file)
--> 326     top = _from_json(json_dict)
    327     return top

File ~/Dropbox/Mac/Documents/Vanderbilt/Research/MoSDeF/gmso/gmso/formats/json.py:184, in _from_json(json_dict)
    182 for atom_dict in json_dict["atoms"]:
    183     atom_type_id = atom_dict.pop("atom_type", None)
--> 184     atom = Atom.model_validate(atom_dict)
    185     top.add_site(atom)
    186     if atom_type_id:

File ~/Dropbox/Mac/Documents/Vanderbilt/Research/MoSDeF/gmso/gmso/abc/gmso_base.py:69, in GMSOBase.model_validate(cls, obj)
     67 @classmethod
     68 def model_validate(cls: Type["Model"], obj: Any) -> "Model":
---> 69     dict_to_unyt(obj)
     70     return super(GMSOBase, cls).model_validate(obj)

File ~/Dropbox/Mac/Documents/Vanderbilt/Research/MoSDeF/gmso/gmso/abc/serialization_utils.py:35, in dict_to_unyt(dict_obj)
     32 else:
     33     unyt_func = u.unyt_array
---> 35 dict_obj[key] = unyt_func(np_array, value["unit"])

File ~/miniconda3/envs/gmso-dev/lib/python3.11/site-packages/unyt/array.py:2231, in unyt_quantity.__new__(cls, input_scalar, units, registry, dtype, bypass_validation, name)
   2229 else:
   2230     units = input_units
-> 2231 ret = unyt_array.__new__(
   2232     cls,
   2233     np.asarray(input_scalar),
   2234     units,
   2235     registry,
   2236     dtype=dtype,
   2237     bypass_validation=bypass_validation,
   2238     name=name,
   2239 )
   2240 if ret.size > 1:
   2241     raise RuntimeError("unyt_quantity instances must be scalars")

File ~/miniconda3/envs/gmso-dev/lib/python3.11/site-packages/unyt/array.py:629, in unyt_array.__new__(cls, input_array, units, registry, dtype, bypass_validation, name)
    624         units = input_units
    625 else:
    626     # units kwarg set, but it's not a Unit object.
    627     # don't handle all the cases here, let the Unit class handle if
    628     # it's a str.
--> 629     units = Unit(input_units, registry=registry)
    631 # Attach the units and name
    632 obj.units = units

File ~/miniconda3/envs/gmso-dev/lib/python3.11/site-packages/unyt/unit_object.py:257, in Unit.__new__(cls, unit_expr, base_value, base_offset, dimensions, registry, latex_repr)
    254         _validate_dimensions(dimensions)
    255 else:
    256     # lookup the unit symbols
--> 257     unit_data = _get_unit_data_from_expr(unit_expr, registry.lut)
    258     base_value = unit_data[0]
    259     dimensions = unit_data[1]

File ~/miniconda3/envs/gmso-dev/lib/python3.11/site-packages/unyt/unit_object.py:970, in _get_unit_data_from_expr(unit_expr, unit_symbol_lut)
    967     return (float(unit_expr), sympy_one)
    969 if isinstance(unit_expr, Symbol):
--> 970     return _lookup_unit_symbol(unit_expr.name, unit_symbol_lut)
    972 if isinstance(unit_expr, Pow):
    973     unit_data = _get_unit_data_from_expr(unit_expr.args[0], unit_symbol_lut)

File ~/miniconda3/envs/gmso-dev/lib/python3.11/site-packages/unyt/unit_registry.py:325, in _lookup_unit_symbol(symbol_str, unit_symbol_lut)
    322     return ret
    324 # no dice
--> 325 raise UnitParseError(
    326     f"Could not find unit symbol '{symbol_str}' in the provided symbols."
    327 )

UnitParseError: Could not find unit symbol 'elementary_charge' in the provided symbols.

[daico007]

I think I know the source of this issue, it's basically because unyt can by default does u.elementary_charge but not u.Unit("elementary_charge"). The workaround is to use the GMSOUnitRegistry when doing the unit parsing (u.Unit("elementary_charge", registry=GMSO_UnitRegistry().reg_)). I am working on the similar issue (local branch, but trying pull the PR in the next day or so) for the XML writer, so I can include the fix for this there.

[CalCraven]

Yep! Got a fix for this and will pull a quick PR.