test_linalg_lstsq_op.py

#   Copyright (c) 2021 PaddlePaddle Authors. 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 unittest
import numpy as np
import paddle
import paddle.fluid as fluid
import paddle.fluid.core as core


class LinalgLstsqTestCase(unittest.TestCase):

    def setUp(self):
        self.devices = ["cpu"]
        self.init_config()
        if core.is_compiled_with_cuda() and self.driver == "gels":
            self.devices.append("gpu:0")
        self.generate_input()
        self.generate_output()
        np.random.seed(2022)

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-15
        self.driver = "gelsd"
        self._input_shape_1 = (5, 4)
        self._input_shape_2 = (5, 3)

    def generate_input(self):
        self._input_data_1 = np.random.random(self._input_shape_1).astype(
            self.dtype)
        self._input_data_2 = np.random.random(self._input_shape_2).astype(
            self.dtype)

    def generate_output(self):
        if len(self._input_shape_1) == 2:
            out = np.linalg.lstsq(self._input_data_1,
                                  self._input_data_2,
                                  rcond=self.rcond)
            self._output_solution = out[0]
            self._output_residuals = out[1]
            self._output_rank = out[2]
            self._output_sg_values = out[3]
        elif len(self._input_shape_1) == 3:
            self._output_solution = []
            self._output_residuals = []
            self._output_rank = []
            self._output_sg_values = []
            for i in range(self._input_shape_1[0]):
                out = np.linalg.lstsq(self._input_data_1[i],
                                      self._input_data_2[i],
                                      rcond=self.rcond)
                self._output_solution.append(out[0])
                self._output_residuals.append(out[1])
                self._output_rank.append(out[2])
                self._output_sg_values.append(out[3])

    def test_eager_dygraph(self):
        paddle.disable_static()
        paddle.fluid.framework._disable_legacy_dygraph()
        for dev in self.devices:
            paddle.set_device(dev)
            place = paddle.CPUPlace() if dev == "cpu" else paddle.CUDAPlace(0)
            x = paddle.to_tensor(self._input_data_1,
                                 place=place,
                                 dtype=self.dtype)
            y = paddle.to_tensor(self._input_data_2,
                                 place=place,
                                 dtype=self.dtype)
            results = paddle.linalg.lstsq(x,
                                          y,
                                          rcond=self.rcond,
                                          driver=self.driver)
            self._result_solution = results[0].numpy()
            self._result_residuals = results[1].numpy()
            self._result_rank = results[2].numpy()
            self._result_sg_values = results[3].numpy()
            self.assert_np_close()

    def test_legacy_dygraph(self):
        paddle.disable_static()
        paddle.fluid.framework._enable_legacy_dygraph()
        for dev in self.devices:
            paddle.set_device(dev)
            place = paddle.CPUPlace() if dev == "cpu" else paddle.CUDAPlace(0)
            x = paddle.to_tensor(self._input_data_1,
                                 place=place,
                                 dtype=self.dtype)
            y = paddle.to_tensor(self._input_data_2,
                                 place=place,
                                 dtype=self.dtype)
            results = paddle.linalg.lstsq(x,
                                          y,
                                          rcond=self.rcond,
                                          driver=self.driver)
            self._result_solution = results[0].numpy()
            self._result_residuals = results[1].numpy()
            self._result_rank = results[2].numpy()
            self._result_sg_values = results[3].numpy()
            self.assert_np_close()

    def test_static(self):
        paddle.enable_static()
        for dev in self.devices:
            paddle.set_device(dev)
            place = fluid.CPUPlace() if dev == "cpu" else fluid.CUDAPlace(0)
            with fluid.program_guard(fluid.Program(), fluid.Program()):
                x = paddle.fluid.data(name="x",
                                      shape=self._input_shape_1,
                                      dtype=self._input_data_1.dtype)
                y = paddle.fluid.data(name="y",
                                      shape=self._input_shape_2,
                                      dtype=self._input_data_2.dtype)
                results = paddle.linalg.lstsq(x,
                                              y,
                                              rcond=self.rcond,
                                              driver=self.driver)
                exe = fluid.Executor(place)
                fetches = exe.run(fluid.default_main_program(),
                                  feed={
                                      "x": self._input_data_1,
                                      "y": self._input_data_2
                                  },
                                  fetch_list=[results])
                self._result_solution = fetches[0]
                self._result_residuals = fetches[1]
                self._result_rank = fetches[2]
                self._result_sg_values = fetches[3]
                self.assert_np_close()

    def assert_np_close(self):
        if len(self._input_shape_1) == 2:
            np.testing.assert_allclose(self._result_solution,
                                       self._output_solution,
                                       rtol=1e-3)
            if self._input_shape_1[-2] > self._input_shape_1[
                    -1] and self._output_rank == self._input_shape_1[-1]:
                np.testing.assert_allclose(self._result_residuals,
                                           self._output_residuals,
                                           rtol=1e-5)
            if self.driver in ("gelsy", "gelsd", "gelss"):
                np.testing.assert_allclose(self._result_rank,
                                           self._output_rank,
                                           rtol=1e-5)
            if self.driver in ("gelsd", "gelss"):
                np.testing.assert_allclose(self._result_sg_values,
                                           self._output_sg_values,
                                           rtol=1e-5)
        else:
            for i in range(len(self._output_solution)):
                np.testing.assert_allclose(self._result_solution[i],
                                           self._output_solution[i],
                                           rtol=1e-3)
                if self._input_shape_1[-2] > self._input_shape_1[
                        -1] and self._output_rank[i] == self._input_shape_1[-1]:
                    np.testing.assert_allclose(self._result_residuals[i],
                                               self._output_residuals[i],
                                               rtol=1e-5)
                if self.driver in ("gelsy", "gelsd", "gelss"):
                    np.testing.assert_allclose(self._result_rank[i],
                                               self._output_rank[i],
                                               rtol=1e-5)
                if self.driver in ("gelsd", "gelss"):
                    np.testing.assert_allclose(self._result_sg_values[i],
                                               self._output_sg_values[i],
                                               rtol=1e-5)


class LinalgLstsqTestCase1(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float32'
        self.rcond = 1e-15
        self.driver = "gels"
        self._input_shape_1 = (9, 9)
        self._input_shape_2 = (9, 5)


class LinalgLstsqTestCase2(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-15
        self.driver = "gels"
        self._input_shape_1 = (5, 10)
        self._input_shape_2 = (5, 8)


class LinalgLstsqTestCase3(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-15
        self.driver = "gels"
        self._input_shape_1 = (10, 7, 3)
        self._input_shape_2 = (10, 7, 6)


class LinalgLstsqTestCaseRcond(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-7
        self.driver = "gelsd"
        self._input_shape_1 = (3, 2)
        self._input_shape_2 = (3, 3)


class LinalgLstsqTestCaseGelsFloat32(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float32'
        self.rcond = None
        self.driver = "gels"
        self._input_shape_1 = (10, 5)
        self._input_shape_2 = (10, 8)


class LinalgLstsqTestCaseGelsFloat64(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float32'
        self.rcond = None
        self.driver = "gels"
        self._input_shape_1 = (3, 2, 8)
        self._input_shape_2 = (3, 2, 15)


class LinalgLstsqTestCaseGelssFloat64(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = None
        self.driver = "gelss"
        self._input_shape_1 = (5, 5)
        self._input_shape_2 = (5, 1)


class LinalgLstsqTestCaseGelsyFloat32(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float32'
        self.rcond = 1e-15
        self.driver = "gelsy"
        self._input_shape_1 = (8, 2)
        self._input_shape_2 = (8, 10)


class LinalgLstsqTestCaseBatch1(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float32'
        self.rcond = 1e-15
        self.driver = "gelss"
        self._input_shape_1 = (2, 3, 10)
        self._input_shape_2 = (2, 3, 4)


class LinalgLstsqTestCaseBatch2(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-15
        self.driver = "gels"
        self._input_shape_1 = (10, 8, 6)
        self._input_shape_2 = (10, 8, 10)


class LinalgLstsqTestCaseLarge1(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-15
        self.driver = "gelsd"
        self._input_shape_1 = (200, 100)
        self._input_shape_2 = (200, 50)


class LinalgLstsqTestCaseLarge2(LinalgLstsqTestCase):

    def init_config(self):
        self.dtype = 'float64'
        self.rcond = 1e-15
        self.driver = "gelss"
        self._input_shape_1 = (50, 600)
        self._input_shape_2 = (50, 300)


if __name__ == '__main__':
    unittest.main()