diff --git a/docs/source/topics/data_types.rst b/docs/source/topics/data_types.rst index a2309d4e66..dd762b88c4 100644 --- a/docs/source/topics/data_types.rst +++ b/docs/source/topics/data_types.rst @@ -284,6 +284,27 @@ The user also has to specify the units for both ``x`` and ``y``: In [4]: xy.set_y(np.array([1, 2, 3, 4]), 'Volume Expansion', '%') + To retrieve the ``x`` values and their metadata, you can use the :py:meth:`~aiida.orm.XyData.get_x` method: + +.. code-block:: ipython + + In [5]: x_name, x_array, x_units = xy.get_x() + +Similarly, to retrieve the ``y`` values and their metadata, use the :py:meth:`~aiida.orm.XyData.get_y` method: + +.. code-block:: ipython + + In [6]: y_values = xy.get_y() + In [7]: for y_name, y_array, y_units in y_values: + ...: print(y_name, y_array, y_units) + +You can also retrieve the names of all ``y`` arrays using :py:meth:`~aiida.orm.XyData.get_y_arraynames`: + +.. code-block:: ipython + + In [8]: y_names = xy.get_y_arraynames() + In [9]: print(y_names) + Note that you can set multiple ``y`` values that correspond to the ``x`` grid. Same as for the :py:class:`~aiida.orm.ArrayData`, the names and shapes of the arrays are stored to the database, the content of the arrays is stored to the repository in the `numpy format `_ (``.npy``). diff --git a/src/aiida/orm/nodes/data/array/xy.py b/src/aiida/orm/nodes/data/array/xy.py index aa994bbd30..910398a511 100644 --- a/src/aiida/orm/nodes/data/array/xy.py +++ b/src/aiida/orm/nodes/data/array/xy.py @@ -9,6 +9,13 @@ """This module defines the classes related to Xy data. That is data that contains collections of y-arrays bound to a single x-array, and the methods to operate on them. + +Example: + xy = XyData() + xy.set_x(np.array([1, 2, 3]), 'x', 'unit_x') + xy.set_y([np.array([1, 2, 3]), np.array([4, 5, 6])], ['y', 'z'], ['unit_y', 'unit_z']) + print(xy.get_arraynames()) # ['x_array', 'y_array_0', 'y_array_1'] + print(xy.get_y()) # [('y', array([1, 2, 3]), 'unit_y'), ('z', array([4, 5, 6]), 'unit_z')] """ from __future__ import annotations @@ -172,10 +179,15 @@ def get_y(self) -> list[tuple[str, 'ndarray', str]]: y_units = self.base.attributes.get('y_units') except (KeyError, AttributeError): raise NotExistent('No y units has been set yet!') - y_arrays = [] + + y_arrays = [self.get_array(f'y_array_{i}') for i in range(len(y_names))] + return list(zip(y_names, y_arrays, y_units)) + + def get_y_arraynames(self) -> list[str]: + """Returns the user-provided names of the y-arrays.""" + try: - for i in range(len(y_names)): - y_arrays += [self.get_array(f'y_array_{i}')] + y_names = self.base.attributes.get('y_names') except (KeyError, AttributeError): - raise NotExistent(f'Could not retrieve array associated with y array {y_names[i]}') - return list(zip(y_names, y_arrays, y_units)) + raise NotExistent('No y names have been set yet!') + return y_names diff --git a/tests/orm/nodes/data/test_xy.py b/tests/orm/nodes/data/test_xy.py index 3842157b78..dd28e4fd1f 100644 --- a/tests/orm/nodes/data/test_xy.py +++ b/tests/orm/nodes/data/test_xy.py @@ -51,3 +51,21 @@ def test_constructor(): node = XyData(x_array, y_array, x_name='x_name', x_units='x_unit', y_names='y_name', y_units='y_units') assert numpy.array_equal(node.get_x()[1], x_array) assert numpy.array_equal(node.get_y()[0][1], y_array) + + +def test_get_y_arraynames(): + """Test retrieving y array names.""" + x_array = numpy.array([1, 2]) + y_array1 = numpy.array([3, 4]) + y_array2 = numpy.array([5, 6]) + + node = XyData() + node.set_x(x_array, 'x_name', 'x_unit') + node.set_y([y_array1, y_array2], ['y_name1', 'y_name2'], ['y_unit1', 'y_unit2']) + + y_names = node.get_y_arraynames() + assert y_names == ['y_name1', 'y_name2'] + + empty_node = XyData() + with pytest.raises(NotExistent): + empty_node.get_y_arraynames()