diff --git a/aiida/orm/nodes/data/array/array.py b/aiida/orm/nodes/data/array/array.py index 06e43a9f43..1bb97e94b3 100644 --- a/aiida/orm/nodes/data/array/array.py +++ b/aiida/orm/nodes/data/array/array.py @@ -193,3 +193,32 @@ def _validate(self): f'Mismatch of files and properties for ArrayData node (pk= {self.pk}): {files} vs. {properties}' ) super()._validate() + + def _get_array_entries(self): + """Return a dictionary with the different array entries. + + The idea is that this dictionary contains the array name as a key and + the value is the numpy array transformed into a list. This is so that + it can be transformed into a json object. + """ + array_dict = {} + for key, val in self.get_iterarrays(): + array_dict[key] = val.tolist() + return array_dict + + def _prepare_json(self, main_file_name='', comments=True): # pylint: disable=unused-argument + """Dump the content of the arrays stored in this node into JSON format. + + :param comments: if True, includes comments (if it makes sense for the given format) + """ + import json + + from aiida import get_file_header + + json_dict = self._get_array_entries() + json_dict['original_uuid'] = self.uuid + + if comments: + json_dict['comments'] = get_file_header(comment_char='') + + return json.dumps(json_dict).encode('utf-8'), {} diff --git a/tests/orm/nodes/data/test_data.py b/tests/orm/nodes/data/test_data.py index 5293ff6444..575ceca35d 100644 --- a/tests/orm/nodes/data/test_data.py +++ b/tests/orm/nodes/data/test_data.py @@ -21,6 +21,7 @@ @pytest.fixture @pytest.mark.usefixtures('aiida_profile_clean') def generate_class_instance(): + # pylint: disable=too-many-return-statements, too-many-statements """Generate a dummy `Data` instance for the given sub class.""" def _generate_class_instance(data_class): @@ -61,6 +62,57 @@ def _generate_class_instance(data_class): instance = data_class(file=filepath_carbon) return instance + if data_class is orm.ArrayData: + instance = data_class() + array_data = numpy.identity(3) + instance.set_array('data', array_data) + instance.set_array('contains_nan_inf', numpy.array([float('NaN'), float('Inf')])) + return instance + + if data_class is orm.KpointsData: + instance = data_class() + cell = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] + instance.set_cell(cell) + instance.set_kpoints_mesh_from_density(0.5) + return instance + + if data_class is orm.XyData: + instance = data_class() + instance.set_x(numpy.arange(5), 'xdata', 'm') + instance.set_y(numpy.arange(5), 'ydata', 'm') + return instance + + if data_class is orm.ProjectionData: + + my_real_hydrogen_dict = { + 'angular_momentum': -3, + 'diffusivity': None, + 'kind_name': 'As', + 'magnetic_number': 0, + 'position': [-1.420047044832945, 1.420047044832945, 1.420047044832945], + 'radial_nodes': 0, + 'spin': 0, + 'spin_orientation': None, + 'x_orientation': None, + 'z_orientation': None + } + kpoints = orm.KpointsData() + kpoints.set_cell([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + kpoints.set_kpoints([[0., 0., 0.]]) + bands = orm.BandsData() + bands.set_kpointsdata(kpoints) + bands.set_bands([[1.0]]) + + RealHydrogen = plugins.OrbitalFactory('core.realhydrogen') # pylint: disable=invalid-name + orbital = RealHydrogen(**my_real_hydrogen_dict) + + instance = data_class() + instance.set_reference_bandsdata(bands) + instance.set_projectiondata( + orbital, list_of_pdos=numpy.asarray([1.0]), list_of_energy=numpy.asarray([1.0]), bands_check=False + ) + return instance + raise RuntimeError( 'no instance generator implemented for class `{}`. If you have added a `_prepare_*` method ' 'for this data class, add a generator of a dummy instance here'.format(data_class)