Source code for openasce.discovery.regression_discovery.trace_expm

#    Copyright 2023 AntGroup CO., Ltd.
#
#    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.

# Some of the code implementation is referred from https://github.com/xunzheng/notears
# Modified by Ant Group in 2023

import numpy as np
import scipy.linalg as slin
import torch


class TraceExpm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        """Forward

        Arguments:
            ctx: the context object used to stash information.
            input: the input tensor
        Returns:
            tensor for output
        """
        E = slin.expm(input.detach().numpy())
        f = np.trace(E)
        E = torch.from_numpy(E)
        ctx.save_for_backward(E)
        return torch.as_tensor(f, dtype=input.dtype)

    @staticmethod
    def backward(ctx, grad_output):
        """Backward

        Arguments:
            ctx: the context object used to retrieve the information.
            grad_output: tensor containing the gradient
        Returns:
            tensor
        """
        (E,) = ctx.saved_tensors
        grad_input = grad_output * E.t()
        return grad_input


trace_expm = TraceExpm.apply