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 index a10dfeb..2cc1bbb 100644 --- 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,24 +17,16 @@ public class GroupRepository extends AnyRepository { private final Multimap groupToProvenance = HashMultimap.create(EXPECTED_KEYS, EXPECTED_VALUES_PER_KEY); private final Map> identity = new HashMap<>(); // members ID -> group IDs - - private final Map>> similarity = new HashMap<>(); // members ID -> group IDs + Map>> 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 mIds) { - Multimap 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 map = memberToGroup.computeIfAbsent(method, k -> new HashMap<>()) + .computeIfAbsent(cutoff, k -> HashMultimap.create()); mIds.forEach(id -> map.put(id, gId)); } @@ -42,22 +34,34 @@ public Collection getGroupToProvenance(String groupId) { return groupToProvenance.get(groupId); } - public Collection 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 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 getMemberToGroup(Input.AggregationMethod method, Integer cutoff, String mId) { + if (isSimilarityBasedGrouping(method)) { + if (cutoff == null) { + List all = new ArrayList<>(); + memberToGroup.get(method).keySet() + .forEach(c -> { + Multimap 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 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(); + } }