ProteinClusterTools 1.0__tar.gz

proteinclustertools-1.0/.gitignore

@@ -0,0 +1,5 @@
+output/
+tmp/
+*.egg-info
+__pycache__
+commands_record_LacI

proteinclustertools-1.0/License

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 John Chen, Barnabas Gall
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

proteinclustertools-1.0/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.3
+Name: ProteinClusterTools
+Version: 1.0
+Summary: Pipeline for clustering protein sequences in large scale datasets, using homology or pLM representation.
+Project-URL: Homepage, https://https://github.com/johnchen93/ProteinClusterTools
+Project-URL: Issues, https://https://github.com/johnchen93/ProteinClusterTools/issues
+Author-email: John Chen <jo.chn.93@gmail.com>, Barnabas Gall <barnabas.gall@anu.edu.au>
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.6

proteinclustertools-1.0/pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "ProteinClusterTools"
+version = "1.0"
+authors = [
+  { name="John Chen", email="jo.chn.93@gmail.com" },
+  { name="Barnabas Gall", email="barnabas.gall@anu.edu.au"}
+]
+description = "Pipeline for clustering protein sequences in large scale datasets, using homology or pLM representation."
+readme = "README.md"
+requires-python = ">=3.6"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+Homepage = "https://https://github.com/johnchen93/ProteinClusterTools"
+Issues = "https://https://github.com/johnchen93/ProteinClusterTools/issues"

proteinclustertools-1.0/src/proteinclustertools/layout/circle_collision.py

@@ -0,0 +1,196 @@
+from dataclasses import dataclass
+import time
+from typing import Optional
+import math
+import Box2D
+import matplotlib.pyplot as plt
+import random
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+
+# Set up the Box2D world
+# world = Box2D.b2World(gravity=(0,0))
+
+@dataclass
+class Circle():
+    x:float
+    y:float
+    r:float
+    id:str
+    parent:Optional['Circle']=None
+    size:Optional[float]=None
+
+# Define a function to create a circle body
+def create_circle(world, x, y, radius, density=1.0):
+    body_def = Box2D.b2BodyDef()
+    body_def.type = Box2D.b2_dynamicBody
+    body_def.position = (x , y )
+    body = world.CreateBody(body_def)
+
+    shape = Box2D.b2CircleShape(radius=radius)
+    fixture_def = Box2D.b2FixtureDef(shape=shape, density=density)
+    body.CreateFixture(fixture_def)
+
+    return body
+
+def GetRadii(sizes, min_size=0, scale=1, force=None):
+    # expect a series of cluster sizes
+    if force is not None:
+        sel=sizes[(sizes>=min_size)|(force)]
+    else:
+        sel=sizes[sizes>=min_size]
+    radii=[]
+    for s in sel:
+        radii.append((s*scale/math.pi)**0.5)
+    radii=pd.Series(radii, index=sel.index, name='radius')
+    return radii
+
+def MakeLayout(radii, radius_pad=1, space_mult=1.2, cycles=10_000, time_step=1, velocity_iterations=6, position_iterations=2,
+                parent:Circle=None, seed=None, pull=0, min_radius_frac=.05, seed_coords=None, expand_mult=1, metadata=None, silence=True):
+    '''
+    radii: pandas Series of radii (index indicates the id of the circle)
+    parent: parent circle to nest current circles within
+    seed: random seed
+    seed_coords: Series of (x,y) coordinates to use as starting points for circles, indexed by the id which is expected to match radii
+    '''
+    if seed is not None:
+        random.seed(seed)
+
+    # if parent is None:
+    area=sum(math.pi*(r+radius_pad)**2 for r in radii)
+    spacing=area**0.5 * space_mult
+
+    world = Box2D.b2World(gravity=(0,0))
+    # bodies=[]
+    rmin, rmax = min(radii), max(radii)
+    min_effective_r = rmax * min_radius_frac
+    total_area=0
+    start_circles=[]
+    N=len(radii)
+    lastx, lasty, lastr=0, 0, 0
+    distances=[]
+    radii_sum=[]
+    for i, r in radii.items():
+        if seed_coords is not None:
+            row=seed_coords.loc[i] # match by index
+            x, y=row['x'], row['y']
+        else:
+            angle=random.random() * 2 * math.pi
+            pos_r=2**-math.log(max(min_effective_r, r)) * spacing + random.random() * spacing * .01
+            x=pos_r*math.cos(angle)
+            y=pos_r*math.sin(angle)
+        total_area+=math.pi*(r+radius_pad)**2
+        # bodies.append(create_circle(world, x, y, r+radius_pad, r**2))
+        start_circles.append(Circle(x, y, r+radius_pad, i))
+        # track distances
+        if i != radii.index[0]:
+            distances.append(((x-lastx)**2+(y-lasty)**2)**0.5)
+            radii_sum.append(r+lastr)
+        lastx, lasty, lastr=x, y, r
+    if len(start_circles)>1:
+        avg_dist=sum(distances)/len(distances)
+        avg_radii_overlap=sum([radii_sum[i]-distances[i] for i in range(len(radii_sum))])/len(radii_sum)
+    else:
+        avg_dist=0
+        avg_radii_overlap=0
+
+    # get boundaries
+    xmax, xmin, ymax, ymin=GetLimits(start_circles)
+    # space the circles out so that the bounding box is 1.2 times the area of the circles
+    if seed_coords is None:
+        expand_factor=1
+        pull_mod=1
+    else:
+        expand_factor=(total_area/(xmax-xmin)/(ymax-ymin)) * expand_mult * len(radii)**.5 * avg_radii_overlap
+        pull_mod=expand_factor/cycles
+    
+    # get centroid of xy
+    # if expand_mult!=1:
+    cx, cy=0, 0
+    for c in start_circles:
+        cx+=c.x
+        cy+=c.y
+    cx/=len(start_circles)
+    cy/=len(start_circles)
+
+    start=time.time()
+    if N >100 and not silence:
+        print(f' Average distance: {avg_dist}, average overlap: {avg_radii_overlap}')
+        print(f' Expanding by {expand_factor} for {len(radii)} circles')
+
+    # scale distance from centroid
+    for c in start_circles:
+        c.x=(c.x-cx)*expand_factor+cx
+        c.y=(c.y-cy)*expand_factor+cy
+
+    if N >100 and not silence:
+        print(f'  Expanding took {time.time()-start} seconds')
+
+    bodies=[]
+    for c in start_circles:
+        bodies.append(create_circle(world, c.x, c.y, c.r, c.r**2))
+
+    for cycle in range(cycles):
+        cycle_mod=2 if cycle<cycles*2/3 else 1
+        if N>1000:
+            cycle_start=time.time()
+        if pull!=0:
+            for c in bodies:
+                x,y=c.position
+                mod=c.fixtures[0].density * pull * pull_mod
+                c.ApplyForceToCenter((-x*mod,-y*mod), True)
+        world.Step(time_step*cycle_mod, velocity_iterations, position_iterations)
+        if N>1000 and not silence:
+            print(f'  Cycle took {time.time()-cycle_start} seconds')
+
+    if N>100 and not silence:
+        print(f' Layout took {time.time()-start} seconds')
+        
+
+    circles=[] 
+    for i in range(len(bodies)):
+        x,y=bodies[i].position
+        id=radii.index[i]
+        circles.append(Circle(x, y, radii[id], id, parent))
+    xmax, xmin, ymax, ymin=GetLimits(circles)
+
+    # get center to off set grouping
+    cx=(xmax+xmin)/2    
+    cy=(ymax+ymin)/2
+
+    xoff, yoff= 0, 0
+    if parent is not None:
+        xoff, yoff=parent.x, parent.y
+    for c in circles:
+        c.x=c.x-cx+xoff
+        c.y=c.y-cy+yoff
+
+    return circles, metadata
+
+def GetLimits(circles, pad=0):
+    xmax, xmin, ymax, ymin=circles[0].x, circles[0].x, circles[0].y, circles[0].y
+    for c in circles:
+        xmax=max(xmax, c.x+c.r)
+        xmin=min(xmin, c.x-c.r)
+        ymax=max(ymax, c.y+c.r)
+        ymin=min(ymin, c.y-c.r)
+    return xmin-pad, xmax+pad, ymin-pad, ymax+pad
+
+if __name__=='__main__':
+
+    # test code
+    n=10_000
+    radii=GetRadii(pd.Series([random.random()*100+1 for x in range(n)]))
+    circles=MakeLayout(radii)
+    xmin, xmax, ymin, ymax=GetLimits(circles)
+
+    f, ax=plt.subplots(figsize=(10,10))
+    ax.axis('off')
+
+    ax.set_xlim(xmin, xmax)
+    ax.set_ylim(ymin, ymax)
+    for c in circles:
+        ax.add_patch(plt.Circle((c.x, c.y), c.r, alpha=.2, linewidth=1, fill=False))
+    plt.show()
+

proteinclustertools-1.0/src/proteinclustertools/layout/circle_layout_tools.py

@@ -0,0 +1,226 @@
+
+import math
+from multiprocessing import Pool
+import os
+from types import SimpleNamespace
+from .circle_collision import Circle, GetRadii, MakeLayout, GetLimits
+import pandas as pd
+import yaml
+
+def ReadConfig(config_file):
+    config_dict=yaml.safe_load(open(config_file))
+    config=SimpleNamespace(**config_dict)
+    return config
+
+def CommentedConfigString(config:SimpleNamespace):
+    cstr=yaml.safe_dump(config.__dict__)
+    commented='#'+cstr.replace('\n','\n#')
+    return commented
+
+def ConfigFromComments(text):
+    cstr=text[1:].replace('\n#','\n')
+    config=SimpleNamespace(**yaml.safe_load(cstr))
+    return config
+
+def process_child_clusters(args):
+    '''Process child clusters in parallel
+        Has same inputs as MakeLayout, but with child_clusters instead of radii as the first argument
+    '''
+    child_clusters, radii, radius_pad, space_scale, seed, seed_pos, cycles, parent, pull, expand_mult, velocity_iterations, position_iterations, time_step, metadata = args
+    return MakeLayout(radii.loc[child_clusters], radius_pad=0, space_mult=space_scale*.8, seed=seed, seed_coords=seed_pos,
+                      time_step=time_step, cycles=cycles, parent=parent, pull=pull, expand_mult=expand_mult,
+                      velocity_iterations=velocity_iterations, position_iterations=position_iterations, metadata=metadata)
+
+def CenterOfMass(circles): # not used
+    x, y=0, 0
+    for c in circles:
+        x+=c.x
+        y+=c.y
+    x/=len(circles)
+    y/=len(circles)
+    return x, y
+
+def Box2DLayout(clusters, levels, minsize=0, space_scale=1.1, pull=1, radius_pad=3, force_inclusion=None, velocity_iterations=16, position_iterations=6,
+                cycle_scale=1000, max_cycles=100_000, time_step=1, pull_base=False, seed=None, seed_coords=None, expand_mult=1):
+    
+    '''
+    Function to turn a table of cluster definitions into a hierarchical layout of circles
+
+    Args:
+        General:
+            clusters (pd.DataFrame): table of cluster definitions, with columns 'id' for each sequence one column per "level" of clustering
+            levels (list): list of column names in clusters to use as levels of clustering. Must be in ascending order (lowest cut-off first) of hierarchy
+            minsize (int): minimum size of a cluster to be included in the layout
+            force_inclusion (list): list of sequence ids to force inclusion in the layout
+        Tuning layout:
+            pull (float): multiplier for the pull between circles when packing
+            radius_pad (float): padding between circles in the base (lowest) level
+            seed (int): random seed for layout. Use same seed to reproduce layout, assuming same input data and cycles of simulation
+            seed_coords (pd.DataFrame): table of seed coordinates for each cluster in the base level. Must have columns 'id', 'x', 'y'
+            pull_base (bool): whether to allow circles in the base level to pull each other. Still respects radius_pad, just makes the base level more compact.
+            Niche:
+                space_scale (float): multiplier for the space between circles on initial randomized layout.
+                expand_mult (float): multiplier for the expansion of circles. Somewhat unpredictable, can increase speed if too many circles start overlapping, but quite situational.
+        Box2D simulation:
+            velocity_iterations (int): number of velocity iterations in the Box2D simulation
+            position_iterations (int): number of position iterations in the Box2D simulation
+            cycle_scale (int): multiplier for the number of cycles in the Box2D simulation per 50 child clusters
+            max_cycles (int): maximum number of cycles in the Box2D simulation
+            time_step (int): time step in the Box2D simulation
+
+    Returns:
+        dict: dictionary of circles at each level of clustering. Each level is a dictionary of circles, with the id of the circle as the key.
+              Each circle object contains its xy position, radius, and size (number of sequences in the cluster)
+    '''
+    
+    ### Process child levels in parallel
+    
+    size_dicts={}
+    if len(levels)>1:
+        process_args = []
+        for lv, level_id in enumerate(levels[1:], start=1):
+            df=clusters.loc[:,['id', level_id]]
+            size=df.groupby(level_id).size().rename('size')#.reset_index()
+            size_dicts[level_id]=dict(size)
+            forced=size.index.isin(df[df['id'].isin(force_inclusion)][level_id].unique()) if force_inclusion is not None else None
+            
+            radii=GetRadii(size, minsize, 1, forced)
+            print(f'Level {level_id}, clusters: {len(radii)}')
+            seed_pos=None
+            if seed_coords is not None:
+                seed_pos=seed_coords.merge(df, on='id')
+                seed_pos=seed_pos[seed_pos[level_id].isin(radii.index)]
+                seed_pos=seed_pos.groupby(level_id)[['x','y']].mean()
+
+            parent_level=levels[lv-1]
+            # subset=clusters[clusters[level_id].isin(radii.index) & clusters[levels[lv-1]].isin(data[parent_level].keys())]
+            subset=clusters[clusters[level_id].isin(radii.index)]
+            
+            print('preparing child clusters')
+            for cluster, group in subset.groupby(parent_level):
+                child_clusters = group[level_id].unique()
+                # print(f'Level {lv} ({level_id}), parent: {cluster}, clusters: {len(group)}, children: {len(child_clusters)}')
+                cycles=min( math.ceil(max(len(child_clusters)/50, 1) * cycle_scale), max_cycles )
+                process_args.append((child_clusters, radii, radius_pad, space_scale, seed, seed_pos,
+                                        cycles, None, pull, expand_mult,
+                                        velocity_iterations, position_iterations, time_step, {'plevel':parent_level, 'parent':cluster,'level':level_id}))    
+
+        print('starting pool')
+        with Pool() as pool:
+            results = list(pool.imap_unordered(process_child_clusters, process_args))
+    else:
+        results=[]
+
+    ### Process base level to avoid collisions
+    level_id=levels[0]
+    df=clusters.loc[:,['id', level_id]]
+    size=df.groupby(level_id).size().rename('size')#.reset_index()
+    size_dicts[level_id]=dict(size)
+    forced=size.index.isin(df[df['id'].isin(force_inclusion)][level_id].unique()) if force_inclusion is not None else None
+    
+    radii=GetRadii(size, minsize, 1, forced)
+    radius_override={}
+    if len(levels)>1:
+        child_level=levels[1]
+        children={}
+        for circles, meta in results:
+            if meta['level']==child_level:
+                children[meta['parent']]=circles
+
+        for cluster, radius in radii.items():
+            if cluster not in children:
+                continue
+            child_circles=children[cluster]
+            # get center of mass
+            # x, y=CenterOfMass(child_circles)
+            x0, x1, y0, y1=GetLimits(child_circles)
+            cx=(x0+x1)/2
+            cy=(y0+y1)/2
+            max_dist=0
+            # distance to each child circle + radius
+            for c in child_circles:
+                dist=((c.x-cx)**2+(c.y-cy)**2)**0.5 + c.r
+                if dist>max_dist:
+                    max_dist=dist
+            if max_dist>radius:
+                radii.loc[cluster]=max_dist
+                radius_override[cluster]=radius
+
+    print(f'Level {level_id}, clusters: {len(radii)}')
+    seed_pos=None
+    if seed_coords is not None:
+        seed_pos=seed_coords.merge(df, on='id')
+        seed_pos=seed_pos[seed_pos[level_id].isin(radii.index)]
+        seed_pos=seed_pos.groupby(level_id)[['x','y']].mean()
+    base_circles, _ =MakeLayout(radii, radius_pad=radius_pad, space_mult=space_scale, seed=seed, 
+                               seed_coords=seed_pos, pull=pull if pull_base else 0, expand_mult=expand_mult,
+                               velocity_iterations=velocity_iterations, position_iterations=position_iterations, time_step=time_step)
+    
+    for c in base_circles:
+        if c.id in radius_override:
+            c.r=radius_override[c.id]
+        c.size=size_dicts[level_id][c.id]
+    # collapse all data
+    data={}
+    data[level_id]=dict((c.id, c) for c in base_circles)
+    grouped={}
+    for circles, meta in results:
+        for c in circles:
+            c.size=size_dicts[meta['level']][c.id]
+        if meta['level'] not in data:
+            grouped[meta['level']]={meta['parent']:{'plevel':meta['plevel'], 'data':circles}}
+            data[meta['level']]=dict((c.id, c) for c in circles)
+        else:
+            grouped[meta['level']][meta['parent']]={'plevel':meta['plevel'], 'data':circles}
+            data[meta['level']].update((c.id, c) for c in circles)
+    # for each level, recenter the circles into the parent circle
+    for lv, level_id in enumerate(levels[1:], start=1):
+        for parent, info in grouped[level_id].items():
+            parent_circle=data[info['plevel']][parent]
+            child_circles=info['data']
+            x0, x1, y0, y1=GetLimits(child_circles)
+            cx=(x0+x1)/2
+            cy=(y0+y1)/2
+
+            px, py=parent_circle.x, parent_circle.y
+            for i, c in enumerate(child_circles):
+                c.x=c.x-cx+px
+                c.y=c.y-cy+py
+    return data
+
+def SaveLayout(data, outfile, config:SimpleNamespace=None):
+    rows=[]
+    for lv, c_dict in data.items():
+        for c in c_dict.values():
+            rows.append([lv, c.id, c.x, c.y, c.r])
+    tb=pd.DataFrame(rows, columns=['level','cluster','x','y','r'])
+    # print(tb)
+    with open(outfile, 'w') as f:
+        if config is not None:
+            print(CommentedConfigString(config), file=f)
+    tb.to_csv(outfile, mode='a', index=False)
+
+def ReadLayout(file, config_check:SimpleNamespace=None):
+    if not os.path.exists(file):
+        return None
+
+    if config_check is not None:
+        saved_cstr=''
+        with open(file) as f:
+            for line in f:
+                if not line.startswith('#'):
+                    break
+                saved_cstr+=line
+        saved_config=ConfigFromComments(saved_cstr)
+        # print(saved_config, config_check)
+        if saved_config!=config_check:
+            return None
+
+    # re-use old layout
+    tb=pd.read_csv(file, comment='#')
+    data={}
+    for lv, group in tb.groupby('level'):
+        data[lv]={}
+        for row in group.itertuples():
+            data[lv][row.cluster]=Circle(row.x,row.y,row.r,row.cluster)
+    return data

proteinclustertools-1.0/src/proteinclustertools/layout/draw_circles.py

@@ -0,0 +1,23 @@
+import matplotlib.pyplot as plt
+
+from layout.circle_collision import GetLimits
+
+def DrawCircles(layout, out_label):
+    '''Debug function to draw a layout of circles to a file'''
+    all_c=[]
+    for lv, c_dict in layout.items():
+        all_c+=list(c_dict.values())
+    xmin, xmax, ymin, ymax=GetLimits(all_c, 5)
+
+    f, ax=plt.subplots(figsize=(10,10))
+    ax.axis('off')
+    ax.set_xlim(xmin, xmax)
+    ax.set_ylim(ymin, ymax)
+    ax.set_aspect('equal', adjustable='box')
+    for lv, c_dict in layout.items():
+        lv_frac=lv/len(layout)
+        for c in c_dict.values():
+            # draw circle
+            circle=plt.Circle((c.x, c.y), c.r, color='b', lw=.5, alpha=.2+lv_frac*.8, fill=False)
+            ax.add_patch(circle)
+    f.savefig(f'{out_label}_layout.png', dpi=300, bbox_inches='tight')