diff --git a/nomad/metainfo/metainfo.py b/nomad/metainfo/metainfo.py index 7aba9c963e151336ba2dd4f00d7286ed28df894e..ef38bb811b77b44ea73c378495225970724e9c9f 100644 --- a/nomad/metainfo/metainfo.py +++ b/nomad/metainfo/metainfo.py @@ -27,6 +27,7 @@ import inspect import re import json import numpy as np +import pandas as pd import pint import pint.unit import pint.quantity @@ -447,6 +448,10 @@ class _QuantityType(DataType): if isinstance(value, np.dtype): return value + if isinstance(value, (pd.core.frame.DataFrame, pd.DataFrame)): + value = self.__to_np(quantity_def, value) + return value + if isinstance(value, Section): return value @@ -1412,6 +1417,14 @@ class MSection(metaclass=MObjectMeta): # TODO find a way to make this a subclas value = value.to(quantity_def.unit).magnitude + if isinstance(value, pd.DataFrame): + try: + value = value.to_numpy() + except TypeError: + raise TypeError( + 'Could not convert value %s of type pandas.Dataframe to a numpy array' % + (value)) + if type(value) != np.ndarray: if len(quantity_def.shape) > 0: try: @@ -1446,6 +1459,8 @@ class MSection(metaclass=MObjectMeta): # TODO find a way to make this a subclas if isinstance(quantity_def.type, np.dtype): value = self.__to_np(quantity_def, value) + elif isinstance(quantity_def.type, pd.DataFrame): + value = self.__to_np(quantity_def, value) else: dimensions = len(quantity_def.shape) if dimensions == 0: diff --git a/tests/metainfo/test_metainfo.py b/tests/metainfo/test_metainfo.py index ae35b059ba649c9913573f7465b01aa2307b0f3f..7a2b46dfa4a1be5bbec2403fbc9306e877d64a8d 100644 --- a/tests/metainfo/test_metainfo.py +++ b/tests/metainfo/test_metainfo.py @@ -21,6 +21,7 @@ import pytest import numpy as np +import pandas as pd import pint.quantity from nomad.metainfo.metainfo import ( @@ -491,6 +492,12 @@ class TestM1: system.atom_positions = [[1, 2, 3]] assert isinstance(system.atom_positions, pint.quantity._Quantity) + def test_pd_dataframe(self): + system = System() + system.atom_positions = pd.DataFrame([[1, 2], [3, 4]]) + assert isinstance(system.atom_positions, pint.quantity._Quantity) + assert np.all(system.atom_positions.m == [[1, 2], [3, 4]]) + def test_np_scalar(self): class TestSection(MSection): test_quantity = Quantity(type=np.dtype('int16')) @@ -500,6 +507,14 @@ class TestM1: assert test_section.test_quantity == 12 assert type(test_section.test_quantity) == np.int16 + def test_pd_dataframe_quantity(self): + class TestSection(MSection): + test_quantity = Quantity(type=np.dtype('int16')) + + test_section = TestSection() + test_section.test_quantity = pd.DataFrame([[1, 2]]) + assert np.all(test_section.test_quantity == [1, 2]) + def test_unit_conversion(self): system = System() system.atom_positions = [[1, 2, 3]] * ureg.angstrom