text: # -*- coding: utf-8 -*-
  ##########################################################################
  # NSAp - Copyright (C) CEA, 2020
  # Distributed under the terms of the CeCILL-B license, as published by
  # the CEA-CNRS-INRIA. Refer to the LICENSE file or to
  # http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
  # for details.
  ##########################################################################
  
  # System import
  import unittest
  import numpy as np
  import torch
  import torch.nn as nn
  
  # Package import
  import pynet
  
  
  class TestLosses(unittest.TestCase):
      """ Test the losses defined in pynet.
      """
      def setUp(self):
          """ Setup test.
          """
          self.losses = pynet.get_tools(tool_name="losses")
          self.n_classes = 3
          self.x = torch.randn(2, self.n_classes, 3, 5, 5, requires_grad=True)
          self.target = torch.empty(2, 3, 5, 5, dtype=torch.long).random_(
              self.n_classes)
          self.weights = torch.tensor([1., 2., 3.])
          self.mask = torch.ones(2, 3, 5, 5)
  
      def tearDown(self):
          """ Run after each test.
          """
          pass
  
      def test_focal(self):
          """ Test the FocalLoss.
          """
          criterion = self.losses["FocalLoss"](
              n_classes=self.n_classes, gamma=0, reduction="mean",
              with_logit=True, alpha=self.weights.numpy().tolist())
          ref_criterion = nn.CrossEntropyLoss(weight=self.weights)
          loss = criterion(self.x, self.target)
          ref_loss = ref_criterion(self.x, self.target)
          alt_loss = criterion._forward_without_resizing(self.x, self.target)
          loss.backward()
          self.assertTrue(np.allclose(
              loss.detach().numpy(), ref_loss.detach().numpy()))
          self.assertTrue(np.allclose(
              loss.detach().numpy(), alt_loss.detach().numpy()))
  
      def test_mask(self):
          """ Test the MaskLoss.
          """
          criterion = self.losses["MaskLoss"](
              n_classes=self.n_classes, beta=1., reduction="mean",
              with_logit=True, alpha=self.weights.numpy().tolist())
          ref_criterion = nn.CrossEntropyLoss(weight=self.weights)
          loss = criterion(self.x, self.target, self.mask)
          ref_loss = ref_criterion(self.x, self.target)
          loss.backward()
          self.assertTrue(np.allclose(
              loss.detach().numpy(), ref_loss.detach().numpy()))
  
      def test_softdice(self):
          """ Test the SoftDiceLoss.
          """
          criterion = self.losses["SoftDiceLoss"](
              reduction="mean", with_logit=True)
          loss = criterion(self.x, self.target)
          alt_loss = criterion._forward_without_resizing(self.x, self.target)
          loss.backward()
          self.assertTrue(np.allclose(
              loss.detach().numpy(), alt_loss.detach().numpy()))
  
      def test_mse(self):
          """ Test the MSELoss.
          """
          criterion = self.losses["MSELoss"]()
          ref_criterion = nn.MSELoss(reduction="mean")
          loss = criterion(self.x, self.x)
          ref_loss = ref_criterion(self.x, self.x)
          loss.backward()
          self.assertTrue(np.allclose(
              loss.detach().numpy(), ref_loss.detach().numpy()))
  
      def test_pcc(self):
          """ Test the PCCLoss.
          """
          criterion = self.losses["PCCLoss"]()
          loss = criterion(self.x, self.x)
          loss.backward()
  
      def test_ncc(self):
          """ Test the NCCLoss.
          """
          criterion = self.losses["NCCLoss"]()
          loss = criterion(self.x[:, :1], self.x[:, :1])
          loss.backward()
          self.assertTrue(np.allclose(np.abs(loss.detach().numpy()), 1))
  
  
  if __name__ == "__main__":
      from pynet.utils import setup_logging
      setup_logging(level="debug")
      unittest.main()
  