# coding=utf-8
# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
#
# 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 numpy as np
import matplotlib.pyplot as plt
[docs]
class GraphDAG(object):
'''
Visualization for causal discovery learning results.
Parameters
----------
est_dag: np.ndarray
The DAG matrix to be estimated.
true_dag: np.ndarray
The true DAG matrix.
show: bool
Select whether to display pictures.
save_name: str
The file name of the image to be saved.
'''
def __init__(self, est_dag, true_dag=None, show=True, save_name=None):
self.est_dag = est_dag
self.true_dag = true_dag
self.show = show
self.save_name = save_name
if not isinstance(est_dag, np.ndarray):
raise TypeError("Input est_dag is not numpy.ndarray!")
if true_dag is not None and not isinstance(true_dag, np.ndarray):
raise TypeError("Input true_dag is not numpy.ndarray!")
if not show and save_name is None:
raise ValueError('Neither display nor save the picture! ' + \
'Please modify the parameter show or save_name.')
GraphDAG._plot_dag(self.est_dag, self.true_dag, self.show, self.save_name)
@staticmethod
def _plot_dag(est_dag, true_dag, show=True, save_name=None):
"""
Plot the estimated DAG and the true DAG.
Parameters
----------
est_dag: np.ndarray
The DAG matrix to be estimated.
true_dag: np.ndarray
The True DAG matrix.
show: bool
Select whether to display pictures.
save_name: str
The file name of the image to be saved.
"""
if isinstance(true_dag, np.ndarray):
# trans diagonal element into 0
for i in range(len(true_dag)):
if est_dag[i][i] == 1:
est_dag[i][i] = 0
if true_dag[i][i] == 1:
true_dag[i][i] = 0
# set plot size
fig, (ax1, ax2) = plt.subplots(figsize=(8, 3), ncols=2)
ax1.set_title('est_graph')
map1 = ax1.imshow(est_dag, cmap='Greys', interpolation='none')
fig.colorbar(map1, ax=ax1)
ax2.set_title('true_graph')
map2 = ax2.imshow(true_dag, cmap='Greys', interpolation='none')
fig.colorbar(map2, ax=ax2)
if save_name is not None:
fig.savefig(save_name)
if show:
plt.show()
elif isinstance(est_dag, np.ndarray):
# trans diagonal element into 0
for i in range(len(est_dag)):
if est_dag[i][i] == 1:
est_dag[i][i] = 0
# set plot size
fig, ax1 = plt.subplots(figsize=(4, 3), ncols=1)
ax1.set_title('est_graph')
map1 = ax1.imshow(est_dag, cmap='Greys', interpolation='none')
fig.colorbar(map1, ax=ax1)
if save_name is not None:
fig.savefig(save_name)
if show:
plt.show()