From 82db9ac0e9141f0c3684b329fbd91c56f0c96995 Mon Sep 17 00:00:00 2001
From: Peter Steinbach <p.steinbach@hzdr.de>
Date: Fri, 27 Nov 2020 17:09:40 +0100
Subject: [PATCH] refactored coverage and density into their own functions

---
 prdc/__init__.py |  2 +-
 prdc/prdc.py     | 52 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/prdc/__init__.py b/prdc/__init__.py
index 859a6ae..d530f8e 100644
--- a/prdc/__init__.py
+++ b/prdc/__init__.py
@@ -1 +1 @@
-from .prdc import compute_prdc
+from .prdc import compute_prdc, compute_coverage, compute_density
diff --git a/prdc/prdc.py b/prdc/prdc.py
index 774681f..d7f3dd0 100644
--- a/prdc/prdc.py
+++ b/prdc/prdc.py
@@ -6,7 +6,7 @@
 import numpy as np
 import sklearn.metrics
 
-__all__ = ['compute_prdc']
+__all__ = ['compute_prdc', 'compute_density', 'compute_coverage']
 
 
 def compute_pairwise_distance(data_x, data_y=None):
@@ -51,6 +51,56 @@ def compute_nearest_neighbour_distances(input_features, nearest_k):
     return radii
 
 
+def compute_density(real_features, fake_features, nearest_k):
+    """
+    Computes density given two manifolds.
+
+    Args:
+        real_features: numpy.ndarray([N, feature_dim], dtype=np.float32)
+        fake_features: numpy.ndarray([N, feature_dim], dtype=np.float32)
+        nearest_k: int.
+    Returns:
+        density: np.float64
+    """
+
+    real_nearest_neighbour_distances = compute_nearest_neighbour_distances(
+        real_features, nearest_k)
+    distance_real_fake = compute_pairwise_distance(
+        real_features, fake_features)
+
+    density = (1. / float(nearest_k)) * (
+            distance_real_fake <
+            np.expand_dims(real_nearest_neighbour_distances, axis=1)
+    ).sum(axis=0).mean()
+
+    return density
+
+
+def compute_coverage(real_features, fake_features, nearest_k):
+    """
+    Computes coverage given two manifolds.
+
+    Args:
+        real_features: numpy.ndarray([N, feature_dim], dtype=np.float32)
+        fake_features: numpy.ndarray([N, feature_dim], dtype=np.float32)
+        nearest_k: int.
+    Returns:
+        coverage: np.float64
+    """
+
+    real_nearest_neighbour_distances = compute_nearest_neighbour_distances(
+        real_features, nearest_k)
+    distance_real_fake = compute_pairwise_distance(
+        real_features, fake_features)
+
+    coverage = (
+            distance_real_fake.min(axis=1) <
+            real_nearest_neighbour_distances
+    ).mean()
+
+    return coverage
+
+
 def compute_prdc(real_features, fake_features, nearest_k):
     """
     Computes precision, recall, density, and coverage given two manifolds.