Repository: skypyproject/skypy

Issue:
set default cosmology from pipeline
When a `cosmology` is set in a pipeline configuration, it should be registered as the default cosmology during the run:

```py
with default_cosmology.set(...):
    ...
```

This probably needs taking the `cosmology` variable out of the normal dependency graph and process it before anything else is run.

Discussion:

Patch:
diff --git a/skypy/pipeline/_pipeline.py b/skypy/pipeline/_pipeline.py
index 01b70a0..c6ab8de 100644
--- a/skypy/pipeline/_pipeline.py
+++ b/skypy/pipeline/_pipeline.py
@@ -4,8 +4,9 @@ This module provides methods to run pipelines of functions with dependencies
 and handle their results.
 """
 
+from astropy.cosmology import default_cosmology
 from astropy.table import Table
-from copy import deepcopy
+from copy import copy, deepcopy
 from importlib import import_module
 import builtins
 import networkx
@@ -108,6 +109,7 @@ class Pipeline:
         # config contains settings for all variables and table initialisation
         # table_config contains settings for all table columns
         self.config = deepcopy(configuration)
+        self.cosmology = self.config.pop('cosmology', {})
         self.parameters = self.config.pop('parameters', {})
         self.table_config = self.config.pop('tables', {})
         default_table = ('astropy.table.Table',)
@@ -128,6 +130,9 @@ class Pipeline:
         functions = {}
         for job in self.parameters:
             self.dag.add_node(job)
+            self.skip_jobs.add(job)
+        self.dag.add_node('cosmology')
+        self.skip_jobs.add('cosmology')
         for job, settings in self.config.items():
             self.dag.add_node(job)
             if isinstance(settings, tuple):
@@ -188,28 +193,35 @@ class Pipeline:
         self.parameters.update(parameters)
 
         # initialise state object
-        self.state = {}
-
-        # go through the jobs in dependency order
-        for job in networkx.topological_sort(self.dag):
-            if job in self.skip_jobs:
-                continue
-            elif job in self.parameters:
-                self.state[job] = self.parameters[job]
-            elif job in self.config:
-                settings = self.config.get(job)
-                self.state[job] = self.get_value(settings)
-            else:
-                table, column = job.split('.')
-                settings = self.table_config[table][column]
-                names = [n.strip() for n in column.split(',')]
-                if len(names) > 1:
-                    # Multi-column assignment
-                    t = Table(self.get_value(settings), names=names)
-                    self.state[table].add_columns(t.columns.values())
+        self.state = copy(self.parameters)
+
+        # Initialise cosmology from config parameters or use astropy default
+        if self.cosmology:
+            self.state['cosmology'] = self.get_value(self.cosmology)
+        else:
+            self.state['cosmology'] = default_cosmology.get()
+
+        # Execute pipeline setting state cosmology as the default
+        with default_cosmology.set(self.state['cosmology']):
+
+            # go through the jobs in dependency order
+            for job in networkx.topological_sort(self.dag):
+                if job in self.skip_jobs:
+                    continue
+                elif job in self.config:
+                    settings = self.config.get(job)
+                    self.state[job] = self.get_value(settings)
                 else:
-                    # Single column assignment
-                    self.state[table][column] = self.get_value(settings)
+                    table, column = job.split('.')
+                    settings = self.table_config[table][column]
+                    names = [n.strip() for n in column.split(',')]
+                    if len(names) > 1:
+                        # Multi-column assignment
+                        t = Table(self.get_value(settings), names=names)
+                        self.state[table].add_columns(t.columns.values())
+                    else:
+                        # Single column assignment
+                        self.state[table][column] = self.get_value(settings)
 
     def write(self, file_format=None, overwrite=False):
         r'''Write pipeline results to disk.


Test patch:
diff --git a/skypy/pipeline/tests/test_pipeline.py b/skypy/pipeline/tests/test_pipeline.py
index 1f9bcd5..450f0b4 100644
--- a/skypy/pipeline/tests/test_pipeline.py
+++ b/skypy/pipeline/tests/test_pipeline.py
@@ -14,6 +14,12 @@ def setup_module(module):
     module.multi_column_array = lambda nrows, ncols: np.ones((nrows, ncols))
     module.multi_column_tuple = lambda nrows, ncols: (np.ones(nrows),) * ncols
 
+    # Define function for testing pipeline cosmology
+    from skypy.utils import uses_default_cosmology
+    @uses_default_cosmology
+    def return_cosmology(cosmology):
+        return cosmology
+    module.return_cosmology = return_cosmology
 
 def test_pipeline():
 
@@ -161,6 +167,35 @@ def test_multi_column_assignment_failure(na, nt):
     with pytest.raises(ValueError):
         pipeline.execute()
 
+def test_pipeline_cosmology():
+
+    # Initial default_cosmology
+    initial_default = default_cosmology.get()
+
+    # Test pipeline correctly sets default cosmology from parameters
+    # N.B. astropy cosmology class has not implemented __eq__ for comparison
+    from astropy.cosmology import FlatLambdaCDM
+    H0, Om0 = 70, 0.3
+    config = {'parameters': {'H0': H0, 'Om0': Om0},
+              'cosmology': ('astropy.cosmology.FlatLambdaCDM', ['$H0', '$Om0']),
+              'test': ('skypy.pipeline.tests.test_pipeline.return_cosmology', ),
+              }
+    pipeline = Pipeline(config)
+    pipeline.execute()
+    assert type(pipeline['test']) == FlatLambdaCDM
+    assert pipeline['test'].H0.value == H0
+    assert pipeline['test'].Om0 == Om0
+
+    # Test pipeline correctly updates cosmology from new parameters
+    H0_new, Om0_new = 75, 0.25
+    pipeline.execute({'H0': H0_new, 'Om0': Om0_new})
+    assert type(pipeline['test']) == FlatLambdaCDM
+    assert pipeline['test'].H0.value == H0_new
+    assert pipeline['test'].Om0 == Om0_new
+
+    # Check that the astropy default cosmology is unchanged
+    assert default_cosmology.get() == initial_default
+
 
 def teardown_module(module):