# coding=utf-8
# 2021.03 modified (1) notears_linear(def) to Notears(class)
# 2021.03 added (1) logging;
# (2) BaseLearner
# 2021.03 deleted (1) __main__
# Huawei Technologies Co., Ltd.
#
# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
#
# Copyright (c) Xun Zheng (https://github.com/xunzheng/notears)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import numpy as np
import scipy.optimize as sopt
from scipy.special import expit as sigmoid
from castle.common import BaseLearner, Tensor
from castle.common.consts import NOTEARS_VALID_PARAMS
from castle.common.validator import check_args_value
[docs]
class Notears(BaseLearner):
"""
Notears Algorithm.
A gradient-based algorithm for linear data models (typically with least-squares loss).
Parameters
----------
lambda1: float
l1 penalty parameter
loss_type: str
l2, logistic, poisson
max_iter: int
max num of dual ascent steps
h_tol: float
exit if |h(w_est)| <= htol
rho_max: float
exit if rho >= rho_max
w_threshold: float
drop edge if |weight| < threshold
Attributes
----------
causal_matrix : numpy.ndarray
Learned causal structure matrix
References
----------
https://arxiv.org/abs/1803.01422
Examples
--------
>>> from castle.algorithms import Notears
>>> from castle.datasets import load_dataset
>>> from castle.common import GraphDAG
>>> from castle.metrics import MetricsDAG
>>> X, true_dag, _ = load_dataset('IID_Test')
>>> n = Notears()
>>> n.learn(X)
>>> GraphDAG(n.causal_matrix, true_dag)
>>> met = MetricsDAG(n.causal_matrix, true_dag)
>>> print(met.metrics)
"""
@check_args_value(NOTEARS_VALID_PARAMS)
def __init__(self, lambda1=0.1,
loss_type='l2',
max_iter=100,
h_tol=1e-8,
rho_max=1e+16,
w_threshold=0.3):
super().__init__()
self.lambda1 = lambda1
self.loss_type = loss_type
self.max_iter = max_iter
self.h_tol = h_tol
self.rho_max = rho_max
self.w_threshold = w_threshold
[docs]
def learn(self, data, columns=None, **kwargs):
"""
Set up and run the Notears algorithm.
Parameters
----------
data: castle.Tensor or numpy.ndarray
The castle.Tensor or numpy.ndarray format data you want to learn.
columns : Index or array-like
Column labels to use for resulting tensor. Will default to
RangeIndex (0, 1, 2, ..., n) if no column labels are provided.
"""
X = Tensor(data, columns=columns)
W_est = self.notears_linear(X, lambda1=self.lambda1,
loss_type=self.loss_type,
max_iter=self.max_iter,
h_tol=self.h_tol,
rho_max=self.rho_max)
causal_matrix = (abs(W_est) > self.w_threshold).astype(int)
self.weight_causal_matrix = Tensor(W_est,
index=X.columns,
columns=X.columns)
self.causal_matrix = Tensor(causal_matrix, index=X.columns,
columns=X.columns)
[docs]
def notears_linear(self, X, lambda1, loss_type, max_iter, h_tol,
rho_max):
"""
Solve min_W L(W; X) + lambda1 ‖W‖_1 s.t. h(W) = 0 using
augmented Lagrangian.
Parameters
----------
X: np.ndarray
n*d sample matrix
Return
------
W_est: np.ndarray
d*d estimated DAG
"""
def _loss(W):
"""Evaluate value and gradient of loss."""
M = X @ W
if loss_type == 'l2':
R = X - M
loss = 0.5 / X.shape[0] * (R ** 2).sum()
G_loss = - 1.0 / X.shape[0] * X.T @ R
elif loss_type == 'logistic':
loss = 1.0 / X.shape[0] * (np.logaddexp(0, M) - X * M).sum()
G_loss = 1.0 / X.shape[0] * X.T @ (sigmoid(M) - X)
elif loss_type == 'poisson':
S = np.exp(M)
loss = 1.0 / X.shape[0] * (S - X * M).sum()
G_loss = 1.0 / X.shape[0] * X.T @ (S - X)
else:
raise ValueError('unknown loss type')
return loss, G_loss
def _h(W):
"""
Evaluate value and gradient of acyclicity constraint.
"""
# E = slin.expm(W * W) # (Zheng et al. 2018)
# h = np.trace(E) - d
M = np.eye(d) + W * W / d # (Yu et al. 2019)
E = np.linalg.matrix_power(M, d - 1)
h = (E.T * M).sum() - d
G_h = E.T * W * 2
return h, G_h
def _adj(w):
"""
Convert doubled variables ([2 d^2] array) back to original
variables ([d, d] matrix).
"""
return (w[:d * d] - w[d * d:]).reshape([d, d])
def _func(w):
"""
Evaluate value and gradient of augmented Lagrangian for
doubled variables ([2 d^2] array).
"""
W = _adj(w)
loss, G_loss = _loss(W)
h, G_h = _h(W)
obj = loss + 0.5 * rho * h * h + alpha * h + lambda1 * w.sum()
G_smooth = G_loss + (rho * h + alpha) * G_h
g_obj = np.concatenate((G_smooth + lambda1, - G_smooth + lambda1),
axis=None)
return obj, g_obj
n, d = X.shape
# double w_est into (w_pos, w_neg)
w_est, rho, alpha, h = np.zeros(2 * d * d), 1.0, 0.0, np.inf
bnds = [(0, 0) if i == j else (0, None) for _ in range(2)
for i in range(d) for j in range(d)]
if loss_type == 'l2':
X = X - np.mean(X, axis=0, keepdims=True)
logging.info('[start]: n={}, d={}, iter_={}, h_={}, rho_={}'.format( \
n, d, max_iter, h_tol, rho_max))
for i in range(max_iter):
w_new, h_new = None, None
while rho < rho_max:
sol = sopt.minimize(_func, w_est, method='L-BFGS-B',
jac=True, bounds=bnds)
w_new = sol.x
h_new, _ = _h(_adj(w_new))
logging.info(
'[iter {}] h={:.3e}, loss={:.3f}, rho={:.1e}'.format( \
i, h_new, _func(w_est)[0], rho))
if h_new > 0.25 * h:
rho *= 10
else:
break
w_est, h = w_new, h_new
alpha += rho * h
if h <= h_tol or rho >= rho_max:
break
W_est = _adj(w_est)
logging.info('FINISHED')
return W_est