def __init__(
self,
vs: Optional[Iterable[str]],
directed_edges: Optional[Iterable[Tuple[str, str]]] = frozenset(),
bidirected_edges: Optional[Iterable[Tuple[str, str,
str]]] = frozenset(),
copy: 'CausalDiagram' = None,
with_do: Optional[Set[str]] = None,
with_induced: Optional[Set[str]] = None):
with_do = wrap(with_do)
with_induced = wrap(with_induced)
if copy is not None:
if with_do is not None:
self.V = copy.V
self.U = wrap(u for u in copy.U
if with_do.isdisjoint(copy.confounded_dict[u]))
self.confounded_dict = {
u: val
for u, val in copy.confounded_dict.items() if u in self.U
}
# copy cautiously
dopa = copy.pa(with_do)
doAn = copy.An(with_do)
doDe = copy.De(with_do)
self._pa = defaultdict(
frozenset, {
k: frozenset() if k in with_do else v
for k, v in copy._pa.items()
})
self._ch = defaultdict(
frozenset, {
k: (v - with_do) if k in dopa else v
for k, v in copy._ch.items()
})
self._an = dict_except(copy._an, doDe)
self._de = dict_except(copy._de, doAn)
elif with_induced is not None:
assert with_induced <= copy.V
removed = copy.V - with_induced
self.V = with_induced
self.confounded_dict = {
u: val
for u, val in copy.confounded_dict.items() if val <= self.V
}
self.U = wrap(self.confounded_dict)
children_are_removed = copy.pa(removed) & self.V
parents_are_removed = copy.ch(removed) & self.V
ancestors_are_removed = copy.de(removed) & self.V
descendants_are_removed = copy.an(removed) & self.V
self._pa = defaultdict(
frozenset, {
x:
(copy._pa[x] -
removed) if x in parents_are_removed else copy._pa[x]
for x in self.V
})
self._ch = defaultdict(
frozenset, {
x:
(copy._ch[x] -
removed) if x in children_are_removed else copy._ch[x]
for x in self.V
})
self._an = dict_only(copy._an, self.V - ancestors_are_removed)
self._de = dict_only(copy._de,
self.V - descendants_are_removed)
else:
self.V = copy.V
self.U = copy.U
self.confounded_dict = copy.confounded_dict
self._ch = copy._ch
self._pa = copy._pa
self._an = copy._an
self._de = copy._de
else:
directed_edges = list(directed_edges)
bidirected_edges = list(bidirected_edges)
self.V = frozenset(vs) | fzset_union(directed_edges) | fzset_union(
(x, y) for x, y, _ in bidirected_edges)
self.U = frozenset(u for _, _, u in bidirected_edges)
self.confounded_dict = {
u: frozenset({x, y})
for x, y, u in bidirected_edges
}
self._ch = pairs2dict(directed_edges)
self._pa = pairs2dict(directed_edges, backward=True)
self._an = dict() # cache
self._de = dict() # cache
assert self._ch.keys() <= self.V and self._pa.keys() <= self.V
self.edges = tuple((x, y) for x, ys in self._ch.items() for y in ys)
self.causal_order = functools.lru_cache()(self.causal_order)
self._do_ = functools.lru_cache()(self._do_)
self.__cc = None
self.__cc_dict = None
self.__h = None
self.__characteristic = None
self.__confoundeds = None
self.u_pas = defaultdict(set)
for u, xy in self.confounded_dict.items():
for v in xy:
self.u_pas[v].add(u)
self.u_pas = defaultdict(
set, {v: frozenset(us)
for v, us in self.u_pas.items()})
def c_component(self, v_or_vs) -> FrozenSet:
assert isinstance(v_or_vs, str)
self.__ensure_cc_cached()
return fzset_union(self.__cc_dict[v] for v in wrap(v_or_vs))
def __de(self, v) -> FrozenSet:
if v in self._de:
return self._de[v]
self._de[v] = fzset_union(self.__de(child)
for child in self._ch[v]) | self._ch[v]
return self._de[v]
def __an(self, v) -> FrozenSet:
if v in self._an:
return self._an[v]
self._an[v] = fzset_union(self.__an(parent)
for parent in self._pa[v]) | self._pa[v]
return self._an[v]
def de(self, v_or_vs) -> FrozenSet:
if isinstance(v_or_vs, str):
return self.__de(v_or_vs)
return fzset_union(self.__de(v) for v in wrap(v_or_vs))
def ch(self, v_or_vs) -> FrozenSet:
if isinstance(v_or_vs, str):
return self._ch[v_or_vs]
else:
return fzset_union(self._ch[v] for v in v_or_vs)