Merge pull request #1100 from giovannipizzi/add_two_band_types

Add two band types in JSON output
aiidateam · Feb 2, 2018 · d88bf93 · d88bf93
2 parents 6dda5f9 + d3cc7e8
commit d88bf93
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/aiida/orm/data/array/bands.py b/aiida/orm/data/array/bands.py
@@ -448,9 +448,11 @@ def _get_bandplot_data(self, cartesian, prettify_format=None, join_symbol=None,
         if len(stored_bands.shape) == 2:
             bands = stored_bands
             band_type_idx = numpy.array([0]*stored_bands.shape[1])
+            two_band_types = False
         elif len(stored_bands.shape) == 3:
             bands = numpy.concatenate([_ for _ in stored_bands], axis=1)
             band_type_idx = numpy.array([0] * stored_bands.shape[2] + [1] * stored_bands.shape[2])
+            two_band_types = True
         else:
             raise ValueError("Unexpected shape of bands")
 
@@ -510,6 +512,7 @@ def _get_bandplot_data(self, cartesian, prettify_format=None, join_symbol=None,
                                  'to': label_to,
                                  'values': bands[position_from:position_to + 1, :].transpose().tolist(),
                                  'x': x[position_from:position_to + 1],
+                                 'two_band_types': two_band_types,
                                  }
                     plot_info['paths'].append(path_dict)
             else:
@@ -520,6 +523,7 @@ def _get_bandplot_data(self, cartesian, prettify_format=None, join_symbol=None,
                              'to': label_to,
                              'values': bands.transpose().tolist(),
                              'x': x,
+                             'two_band_types': two_band_types,
                              }
                 plot_info['paths'].append(path_dict)
                 plot_info['path'].append([label_from, label_to])