allow retrieving all similarity cutoffs

rcsb · Nov 7, 2024 · bf48ba7 · bf48ba7
1 parent 43d4f0d
commit bf48ba7
Showing 1 changed file with 29 additions and 25 deletions.
diff --git a/src/main/java/org/rcsb/idmapper/backend/data/repository/GroupRepository.java b/src/main/java/org/rcsb/idmapper/backend/data/repository/GroupRepository.java
@@ -17,47 +17,51 @@ public class GroupRepository extends AnyRepository {
     private final Multimap<String, String> groupToProvenance = HashMultimap.create(EXPECTED_KEYS, EXPECTED_VALUES_PER_KEY);
 
     private final Map<Input.AggregationMethod,
-            Multimap<String, String>> identity = new HashMap<>(); // members ID -> group IDs
-
-    private final Map<Input.AggregationMethod,
-            Map<Integer, // similarity cutoff
-            Multimap<String, String>>> similarity = new HashMap<>(); // members ID -> group IDs
+            Map<Integer, // similarity cutoff (can be NULL for identity-based groups)
+            Multimap<String, String>>> memberToGroup = new HashMap<>(); // members ID -> group IDs
 
     public void addGroupProvenance(String groupId, String provenanceId) {
         groupToProvenance.put(groupId, provenanceId);
     }
 
     public void addGroupMembers(Input.AggregationMethod method, Integer cutoff, String gId, List<String> mIds) {
-        Multimap<String, String> map;
-        if (cutoff != null) {
-            map = similarity.computeIfAbsent(method, k -> new HashMap<>())
-                    .computeIfAbsent(cutoff, k -> HashMultimap.create());
-        } else {
-            map = identity.computeIfAbsent(method, k -> HashMultimap.create());
-        }
+        Multimap<String, String> map = memberToGroup.computeIfAbsent(method, k -> new HashMap<>())
+                .computeIfAbsent(cutoff, k -> HashMultimap.create());
         mIds.forEach(id -> map.put(id, gId));
     }
 
     public Collection<String> getGroupToProvenance(String groupId) {
         return groupToProvenance.get(groupId);
     }
 
-    public Collection<String> getMemberToGroup(Input.AggregationMethod method, Integer cutoff, String mId) {
-        if (cutoff != null)
-            return similarity.get(method).get(cutoff).get(mId);
-        else
-            return identity.get(method).asMap().getOrDefault(mId, new ArrayList<>());
+    private boolean isSimilarityBasedGrouping(Input.AggregationMethod method) {
+        return method == Input.AggregationMethod.sequence_identity;
     }
 
-    public Long countGroups(Input.AggregationMethod method) {
-        if (method == Input.AggregationMethod.sequence_identity) {
-            AtomicReference<Long> count = new AtomicReference<>(0L);
-            similarity.get(method).keySet()
-                    .forEach(cutoff -> count.updateAndGet(v -> v + similarity.get(method).get(cutoff)
-                            .values().stream().distinct().count()));
-            return count.get();
+    public Collection<String> getMemberToGroup(Input.AggregationMethod method, Integer cutoff, String mId) {
+        if (isSimilarityBasedGrouping(method)) {
+            if (cutoff == null) {
+                List<String> all = new ArrayList<>();
+                memberToGroup.get(method).keySet()
+                        .forEach(c -> {
+                            Multimap<String, String> map = memberToGroup.get(method).get(c);
+                            if (map.containsKey(mId)) {
+                                all.addAll(map.get(mId));
+                            }
+                        });
+                return all;
+            }
+            else return memberToGroup.get(method).get(cutoff).get(mId);
         } else {
-            return identity.get(method).values().stream().distinct().count();
+            return memberToGroup.get(method).get(null).get(mId);
         }
     }
+
+    public Long countGroups(Input.AggregationMethod method) {
+        AtomicReference<Long> count = new AtomicReference<>(0L);
+        memberToGroup.get(method).keySet()
+                .forEach(cutoff -> count.updateAndGet(v -> v + memberToGroup.get(method).get(cutoff)
+                        .values().stream().distinct().count()));
+        return count.get();
+    }
 }