Commits

Chris Mutel  committed 8d86ca5

Only include processes in geomapping array

  • Participants
  • Parent commits 9dd42cc

Comments (0)

Files changed (2)

File bw2data/database.py

         gl = config.global_location
 
         # Create geomapping array
+        count = 0
         arr = np.zeros((len(data), ), dtype=self.dtype_fields_geomapping + self.base_uncertainty_fields)
-        for index, key in enumerate(sorted(data.keys(), key=lambda x: x[1])):
-            arr[index] = (
-                mapping[key],
-                geomapping[data[key].get("location", gl) or gl],
-                MAX_INT_32, MAX_INT_32,
-                0, 1, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, False
-            )
+        for key in sorted(data.keys(), key=lambda x: x[1]):
+            if data[key].get('type') == 'process':
+                arr[count] = (
+                    mapping[key],
+                    geomapping[data[key].get("location", gl) or gl],
+                    MAX_INT_32, MAX_INT_32,
+                    0, 1, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, False
+                )
+                count += 1
 
         filepath = os.path.join(
             config.dir,
             self.name + u".geomapping.pickle"
         )
         with open(filepath, "wb") as f:
-            pickle.dump(arr, f, protocol=pickle.HIGHEST_PROTOCOL)
+            pickle.dump(arr[:count], f, protocol=pickle.HIGHEST_PROTOCOL)
 
         arr = np.zeros((num_exchanges + len(data), ), dtype=self.dtype)
         count = 0

File bw2data/tests/database.py

         fieldnames = {'activity', 'geo', 'row', 'col'}
         self.assertFalse(fieldnames.difference(set(array.dtype.names)))
 
+    def test_only_processes_in_geomapping(self):
+        database = Database("a database")
+        database.register()
+        database.write({
+            ("a database", "foo"): {
+                'exchanges': [],
+                'type': 'process'
+            },
+            ("a database", "bar"): {
+                'exchanges': [],
+                'type': 'process'
+            },
+            ("a database", "baz"): {
+                'exchanges': [],
+                'type': 'not a process'
+            },
+        })
+        database.process()
+        fp = os.path.join(
+            config.dir,
+            u"processed",
+            database.name + u".geomapping.pickle"
+        )
+        array = pickle.load(open(fp, "rb"))
+        self.assertEqual(array.shape, (2,))
+
     def test_geomapping_array_includes_only_processes(self):
         database = Database("a database")
         database.register()
             [x[0] for x in database.dtype_fields],
             ['input', 'output', 'row', 'col', 'type']
         )
-