Make distributions and graph talk to each other #30

mscroggs · 2025-03-20T11:14:10Z

Builds on top of #16.

Removes placeholder distribution classes in graph, and plugs in the proper distribution classes in their place

Co-authored-by: Will Graham <[email protected]>

…nto mscroggs/normal-example

Co-authored-by: Will Graham <[email protected]>

…nto mscroggs/normal-example

…error is nicer

mscroggs · 2025-03-20T11:16:10Z

src/causalprog/algorithms/expectation.py

+    outcome_node_label: str,
+    samples: int,
+    *,
+    rng_key: jax.Array,


mscroggs · 2025-03-20T11:17:27Z

src/causalprog/graph/node.py

+            concrete_dist = self._dist.construct(
+                **parameters, **self._constant_parameters
+            )
+            output[sample] = concrete_dist.sample(new_key[sample], 1)[0][0]


mscroggs · 2025-03-20T11:19:27Z

src/causalprog/graph/node.py

+            return concrete_dist.sample(rng_key, samples)
+        output = np.zeros(samples)
+        new_key = jax.random.split(rng_key, samples)
+        for sample in range(samples):


mscroggs · 2025-03-20T11:24:40Z

src/causalprog/graph/node.py

+        new_key = jax.random.split(rng_key, samples)
+        for sample in range(samples):
+            parameters = {
+                i: sampled_dependencies[j][sample] for i, j in self._parameters.items()


mscroggs · 2025-03-20T15:03:29Z

src/causalprog/distribution/normal.py

@@ -76,7 +78,7 @@ def __init__(self) -> None:
        """Create a family of normal distributions."""
        super().__init__(Normal, family_name="Normal")

-    def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:
+    def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:  # type: ignore # noqa: PGH003


mypy was very unhappy with everything I tried putting here. @willGraham01: Any idea what it should be?

Also, we originally had this as the following, but either ruff or mypy was not happy that the base class didn't have the *,, but having *, **parameters in the base class is a syntax error...

Suggested change

def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal: # type: ignore # noqa: PGH003

def construct(self, *, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal: # type: ignore # noqa: PGH003

I think this is mypy expects any overwritten methods to retain a signature that is compatible with the signature in the base class. Personally I would have also thought that *, mean, cov was compatible with **parameters (since they're all keyword arguments) but in hindsight I can see why it's not (since **parameters does not necessitate ANY arguments, but *, mean, cov necessitates 2).

I think (for now), leave things as they are (*, mean, cov). I can tackle this when I re-work the DistributionFamilys using the backend-agnostic classes (which should make this awkward hand-me-down of parameters redundant).

willGraham01

Looks like a few changes from #16 are being re-overwritten here, but otherwise just a couple of shortenings.

With regards to the construct problems, I'll tackle them in my backend-agnostic PR anyway so happy to postpone with the # noqa until then.

src/causalprog/_abc/labelled.py

willGraham01 · 2025-03-21T10:29:59Z

src/causalprog/distribution/normal.py

@@ -76,7 +78,7 @@ def __init__(self) -> None:
        """Create a family of normal distributions."""
        super().__init__(Normal, family_name="Normal")

-    def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:
+    def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:  # type: ignore # noqa: PGH003


I think this is mypy expects any overwritten methods to retain a signature that is compatible with the signature in the base class. Personally I would have also thought that *, mean, cov was compatible with **parameters (since they're all keyword arguments) but in hindsight I can see why it's not (since **parameters does not necessitate ANY arguments, but *, mean, cov necessitates 2).

I think (for now), leave things as they are (*, mean, cov). I can tackle this when I re-work the DistributionFamilys using the backend-agnostic classes (which should make this awkward hand-me-down of parameters redundant).

src/causalprog/graph/graph.py

src/causalprog/graph/node.py

willGraham01 · 2025-03-21T10:35:08Z

tests/test_distributions/test_family.py

@@ -2,7 +2,7 @@
 import pytest

 from causalprog.distribution.base import SampleTranslator
-from causalprog.distribution.family import DistributionFamily
+from causalprog.distribution.normal import DistributionFamily


Suggested change

from causalprog.distribution.normal import DistributionFamily

from causalprog.distribution.family import DistributionFamily

Not sure why it's now importing from a different module.

I think it's nicer to import from family as that's where this is defined?

tests/test_graph.py

Co-authored-by: Will Graham <[email protected]>

mscroggs and others added 24 commits March 13, 2025 15:23

start working on node that store distributions

03a751f

refactor graph, add add_node and add_edge methods

ddb7d52

get multiple samples at onde

e8a4592

parametrise test

d7c5a60

make two distribution example

bd50473

Merge branch 'main' into mscroggs/normal-example

02695eb

Merge branch 'main' into mscroggs/normal-example

e78a06c

| None

bbe9284

remove irrelevant nodes from test

0a9a7b2

add stdev to test

711ca46

Update src/causalprog/graph/graph.py

e460f4f

Co-authored-by: Will Graham <[email protected]>

Update src/causalprog/graph/graph.py

ef667db

Co-authored-by: Will Graham <[email protected]>

make roots_down_to_outcome a method of the graph

87fa819

Merge branch 'mscroggs/normal-example' of github.com:UCL/causalprog i…

df339e7

…nto mscroggs/normal-example

Update src/causalprog/graph/graph.py

4b457fc

Co-authored-by: Will Graham <[email protected]>

default false

b222449

Merge branch 'mscroggs/normal-example' of github.com:UCL/causalprog i…

ac5dbdc

…nto mscroggs/normal-example

don't allow temporary None labels

fa7bfe7

simpler test

ed93108

improve tests, and simplify iteration code

ed02c31

reduce number of tests

cc2f1ec

number of samples must be int, don't use 0 for mean so that relative …

8c01b7d

…error is nicer

ruff

06a3946

add moved file

8c7a517

mscroggs marked this pull request as draft March 20, 2025 11:14

mscroggs changed the base branch from main to mscroggs/normal-example March 20, 2025 11:14

This was referenced Mar 20, 2025

Remove jax hardcoding #24

Open

[0][0] in return of node.sample #25

Open

remove for loop over samples #26

Open

Make two node example graph #16

Merged

mscroggs commented Mar 20, 2025

View reviewed changes

Base automatically changed from mscroggs/normal-example to main March 20, 2025 14:09

mscroggs added 2 commits March 20, 2025 14:37

Merge branch 'main' into mscroggs/distributions

c519c6e

fix merge

5651633

mscroggs marked this pull request as ready for review March 20, 2025 14:58

mscroggs requested a review from willGraham01 March 20, 2025 14:58

mscroggs commented Mar 20, 2025

View reviewed changes

willGraham01 approved these changes Mar 21, 2025

View reviewed changes

mscroggs and others added 7 commits March 21, 2025 16:49

Update src/causalprog/_abc/labelled.py

6fb79be

Co-authored-by: Will Graham <[email protected]>

Update src/causalprog/graph/graph.py

969f0fe

Co-authored-by: Will Graham <[email protected]>

Update src/causalprog/graph/node.py

66d9317

Co-authored-by: Will Graham <[email protected]>

Update src/causalprog/graph/node.py

ee6cfef

Co-authored-by: Will Graham <[email protected]>

Update tests/test_graph.py

caba090

Co-authored-by: Will Graham <[email protected]>

Merge branch 'main' into mscroggs/distributions

b45ce2c

noqa

f6e69f9

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make distributions and graph talk to each other #30

Make distributions and graph talk to each other #30

mscroggs commented Mar 20, 2025 •

edited

Loading

mscroggs Mar 20, 2025

mscroggs Mar 20, 2025

mscroggs Mar 20, 2025

mscroggs Mar 20, 2025

mscroggs Mar 20, 2025

mscroggs Mar 20, 2025

willGraham01 Mar 21, 2025

willGraham01 left a comment

willGraham01 Mar 21, 2025

willGraham01 Mar 21, 2025

mscroggs Mar 21, 2025

	def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal: # type: ignore # noqa: PGH003
	def construct(self, *, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal: # type: ignore # noqa: PGH003

	from causalprog.distribution.normal import DistributionFamily
	from causalprog.distribution.family import DistributionFamily

Make distributions and graph talk to each other #30

Are you sure you want to change the base?

Make distributions and graph talk to each other #30

Conversation

mscroggs commented Mar 20, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

willGraham01 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mscroggs commented Mar 20, 2025 •

edited

Loading