Refactor of batch data handling and passing to model (#34)

* Refactored handling of data from data loader to model. Functional although final convergence tests are missing. Some cleanup and adding of comments is also missing.

* Fixed problems that surfaces when training with a large number of
streams.

* - Added docs.
- Fixed minor isses, in particular related to evaluation.

* - Added proper docs to StreamData
- Cleaned up some data that was no longer used or could be computed more
  locally.

* - Updated batchifyer (missed in last commit).
- Fixed problem in running evaluation with writing output
- Auto-formatting and linting

Author: clessig

src/weathergen/datasets/batchifyer.py

@@ -314,7 +314,6 @@ def batchify_target(
         if len(source) < 2:
             target_tokens, target_coords = torch.tensor([]), torch.tensor([])
             target_tokens_lens = torch.zeros([self.num_healpix_cells_target], dtype=torch.int32)
-            target_coords_lens = torch.zeros([self.num_healpix_cells_target], dtype=torch.int32)
 
         else:
             thetas = ((90.0 - source[:, geoinfo_offset]) / 180.0) * np.pi
@@ -339,7 +338,6 @@ def batchify_target(
                 target_coords[c] = normalize_targets(t[:, :geoinfo_size].clone())
 
             target_tokens_lens = torch.tensor([len(s) for s in target_tokens], dtype=torch.int32)
-            target_coords_lens = target_tokens_lens.detach().clone()
 
             # if target_coords_local and target_tokens_lens.sum()>0 :
             if target_tokens_lens.sum() > 0:
@@ -352,6 +350,6 @@ def batchify_target(
                     self.hpy_nctrs_target,
                 )
                 target_coords.requires_grad = False
-                target_coords = list(target_coords.split(target_coords_lens.tolist()))
+                target_coords = list(target_coords.split(target_tokens_lens.tolist()))
 
-        return (target_tokens, target_tokens_lens, target_coords, target_coords_lens)
+        return (target_tokens, target_coords)