poseigen-oyster 0.0.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,7 @@
1
+ Copyright 2025 Husam Abdulnabi
2
+
3
+ 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:
4
+
5
+ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
6
+
7
+ 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.
@@ -0,0 +1,25 @@
1
+ Metadata-Version: 2.3
2
+ Name: poseigen_oyster
3
+ Version: 0.0.4
4
+ Summary: The basics for Poseigen software
5
+ License: MIT
6
+ Author: Husam Abdulnabi
7
+ Author-email: husam.abdulnabi@gmail.com
8
+ Maintainer: Husam Abdulnabi
9
+ Maintainer-email: husam.abdulnabi@gmail.com
10
+ Requires-Python: >=3.12
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Topic :: Software Development
14
+ Classifier: Programming Language :: Python :: 3
15
+ Requires-Dist: kneed
16
+ Requires-Dist: numpy
17
+ Requires-Dist: poseigen_chisel
18
+ Requires-Dist: poseigen_seaside
19
+ Requires-Dist: poseigen_trident
20
+ Requires-Dist: scikit-learn
21
+ Requires-Dist: torch
22
+ Project-URL: Homepage, https://github.com/husam94/poseigen_oyster
23
+ Description-Content-Type: text/markdown
24
+
25
+ Will do
@@ -0,0 +1 @@
1
+ Will do
@@ -0,0 +1,65 @@
1
+
2
+
3
+ [build-system]
4
+
5
+ requires = ["setuptools",
6
+ "numpy",
7
+ "torch",
8
+ "poseigen_seaside",
9
+ "poseigen_trident",
10
+ "poseigen_chisel",
11
+ "scikit-learn",
12
+ "kneed"
13
+ ]
14
+
15
+ [project]
16
+
17
+ name = "poseigen_oyster"
18
+
19
+ version = "0.0.4"
20
+
21
+ requires-python = ">=3.12"
22
+
23
+ description = "The basics for Poseigen software"
24
+
25
+ readme = "README.md"
26
+
27
+ license = "MIT"
28
+
29
+ license-files = ["LICEN[CS]E.*"]
30
+
31
+ keywords = []
32
+
33
+ authors = [{ name = "Husam Abdulnabi", email = "husam.abdulnabi@gmail.com" }]
34
+
35
+ maintainers = [
36
+ { name = "Husam Abdulnabi", email = "husam.abdulnabi@gmail.com" },
37
+ ]
38
+
39
+ classifiers = [
40
+ # How mature is this project? Common values are
41
+ # 3 - Alpha
42
+ # 4 - Beta
43
+ # 5 - Production/Stable
44
+ "Development Status :: 3 - Alpha",
45
+
46
+ # Indicate who your project is intended for
47
+ "Intended Audience :: Developers",
48
+ "Topic :: Software Development",
49
+
50
+ "Programming Language :: Python :: 3",
51
+
52
+ ]
53
+
54
+ dependencies = [
55
+ "numpy",
56
+ "torch",
57
+ "poseigen_seaside",
58
+ "poseigen_trident",
59
+ "poseigen_chisel",
60
+ "scikit-learn",
61
+ "kneed"
62
+ ]
63
+
64
+ [project.urls]
65
+ "Homepage" = "https://github.com/husam94/poseigen_oyster"
@@ -0,0 +1,2 @@
1
+ from .model import *
2
+ from .analysis import *
@@ -0,0 +1,451 @@
1
+
2
+ # These functions are specific to EXACT OYSTER
3
+
4
+ import numpy as np
5
+
6
+ import poseigen_chisel as chis
7
+ import poseigen_trident.utils as tu
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12
+
13
+ from sklearn.cluster import KMeans, AgglomerativeClustering
14
+ from kneed import KneeLocator
15
+
16
+
17
+ def SigContrib(model, inp,
18
+ sub_model = None,
19
+ batchsize = 256):
20
+
21
+ shapo = inp.shape #the 2nd dimension is always the number of sigs.
22
+ num_sigs = shapo[1]
23
+
24
+ if isinstance(model, list) is False: model = [model]
25
+ #if sub_model is None: sub_model = ''
26
+
27
+ lk = len(inp)
28
+
29
+ fullbatches = lk // batchsize
30
+ rem = lk % batchsize
31
+
32
+ fins = []
33
+
34
+ for fb in np.arange(fullbatches + (rem > 0)):
35
+
36
+ fo = fb*batchsize
37
+ batch = inp[fo:fo + batchsize]
38
+
39
+ fin = []
40
+
41
+ for mo in model:
42
+
43
+ if isinstance(mo, str): mo = tu.LoadTorch(mo)
44
+
45
+ if sub_model is not None: mo = getattr(mo, sub_model)
46
+
47
+ with torch.no_grad():
48
+ mo.eval()
49
+
50
+ ger = 0 if isinstance(mo.O[0], nn.Conv2d) else 1
51
+
52
+ Dweight = mo.O[ger].weight
53
+ Dweight = torch.unsqueeze(Dweight, 0)
54
+
55
+ x = torch.FloatTensor(batch).to(device)
56
+
57
+ x = mo.Reflect(x)
58
+ x = mo.kE(x)
59
+ x = mo.AntiReflect(x)
60
+ x = mo.P(x)
61
+
62
+ x = torch.unsqueeze(x, axis = 1)
63
+ x = x * Dweight
64
+ x = torch.squeeze(torch.sum(x, axis = 3))
65
+
66
+ xshapo = x.shape
67
+ x = torch.sum(x.reshape(xshapo[0], num_sigs, xshapo[-1] // num_sigs), -1)
68
+
69
+ fin.append(x.cpu().detach().numpy())
70
+
71
+ fin = np.mean(np.stack(fin), axis = 0)
72
+ fins.append(fin)
73
+
74
+ return np.vstack(fins)
75
+
76
+
77
+
78
+ def kmer_contrib(model, inp, sub_model = None,
79
+ joint = False, chunk = False):
80
+
81
+ # The inp is (kmer, *kmer shape)
82
+
83
+ mo = tu.LoadTorch(model) if isinstance(model, str) else model
84
+
85
+ if sub_model is not None: mo = getattr(mo, sub_model)
86
+
87
+ shapo = inp.shape
88
+
89
+ with torch.no_grad():
90
+ mo.eval()
91
+
92
+ ger = 0 if isinstance(mo.O[0], nn.Conv2d) else 1
93
+
94
+ Dweight = mo.O[ger].weight
95
+ Dweight = torch.unsqueeze(Dweight, 0)
96
+
97
+ x = inp
98
+
99
+ if joint: x = mo.Reflect(x)
100
+
101
+ x = mo.kE(x)
102
+ x = torch.unsqueeze(x, axis = 1)
103
+
104
+ #---------------------------
105
+
106
+ if isinstance(mo.P, nn.Identity) is False:
107
+ P_ck = mo.P.kernel_size[0]
108
+ x = x / P_ck
109
+
110
+ #---------------------------
111
+
112
+ x = x * Dweight
113
+
114
+ if chunk:
115
+ xshapo = x.shape
116
+ x = x.reshape(xshapo[0], xshapo[1], shapo[1], xshapo[2] // shapo[1], *xshapo[-2:])
117
+ x = torch.sum(x, axis = 3)
118
+ else:
119
+ x = torch.sum(x, axis = -3)
120
+
121
+ if joint: x = torch.unsqueeze(torch.mean(x, axis = -1), dim= -1)
122
+
123
+ return x.cpu().detach().numpy()
124
+
125
+
126
+ def kmer_contrib_multi(model, inp, sub_model = None,
127
+ joint = False, chunk = False, batchsize = 256):
128
+
129
+ lk = len(inp)
130
+ fullbatches = lk // batchsize
131
+ rem = lk % batchsize
132
+
133
+ kc_args = {'sub_model': sub_model, 'joint': joint, 'chunk': chunk}
134
+
135
+ fins = []
136
+
137
+ for fb in np.arange(fullbatches + (rem > 0)):
138
+
139
+ fo = fb*batchsize
140
+
141
+ batch = inp[fo:fo + batchsize]
142
+
143
+ batch_t = torch.FloatTensor(batch).to(device)
144
+
145
+ if isinstance(model, list):
146
+ fin = [kmer_contrib(m, batch_t, ** kc_args) for m in model]
147
+ fin = np.mean(np.stack(fin), axis = 0)
148
+ else:
149
+ fin = kmer_contrib(model, batch_t, ** kc_args)
150
+
151
+ fins.append(fin)
152
+
153
+ fins = np.vstack(fins)
154
+
155
+ return fins
156
+
157
+
158
+
159
+ def kppm_contrib(model, kppm,
160
+
161
+ sub_model = None, joint = False, chunk = False,
162
+ score_mode = 'recipeuc',
163
+
164
+ random_seqs = 10000, unique = True,
165
+ scoreweight = True, batchsize = 256):
166
+
167
+ # generates kmers from the ppm and scores them to the final output
168
+ # ppm must be the same size as the model head
169
+
170
+ kmers = chis.kmersFromPPM(kppm, num = random_seqs, unique = unique)
171
+ kmers_OHE = np.expand_dims(chis.Seqs2OHE(kmers), axis = 1)
172
+
173
+ fins = kmer_contrib_multi(model, kmers_OHE, sub_model = sub_model,
174
+ joint = joint, chunk = chunk, batchsize = batchsize)
175
+
176
+ if scoreweight:
177
+
178
+ kmers_scored = chis.PWMScorer(np.squeeze(kmers_OHE, axis = 1),
179
+ kppm, revcomp = joint,
180
+ score_mode = score_mode).reshape(-1, 1, 1, 1)
181
+
182
+ out = np.sum((fins * kmers_scored), axis = 0) / np.sum(kmers_scored)
183
+
184
+ else: out = np.mean(fins, axis = 0)
185
+
186
+ return out
187
+
188
+ def contrib_perpos(contribs, oys_dict, submodel = None):
189
+
190
+ # this function takes your contribution per position block
191
+ # and expands it to the whole sequence
192
+
193
+ # center in this case means that the model used padding for pooling.
194
+
195
+ #contribs shape: (subjects, pos, *)
196
+
197
+ #-------------------------------
198
+
199
+ submodel = '' if submodel is None else submodel + '_'
200
+
201
+ seq_len = oys_dict[submodel + 'dim_i'][1]
202
+ k_len, P_ck = [oys_dict[submodel + xoo] for xoo in ['kE_k', 'P_ck']]
203
+ P_center = oys_dict['P_center']
204
+
205
+ #-------------------------------
206
+
207
+ dim_kE_length = seq_len - k_len + 1
208
+
209
+ needed = P_ck - (dim_kE_length % P_ck)
210
+ pad2add = (needed // 2) if P_center else 0
211
+
212
+ P_length = int(np.ceil(dim_kE_length / P_ck))
213
+
214
+ #-------------------------------------------
215
+
216
+ firstpos = P_ck - pad2add
217
+ midpos = P_ck
218
+ lastpos = dim_kE_length - firstpos - (midpos * (P_length - 2))
219
+
220
+ contribs_perpos = [np.repeat(contribs[:, [0]], firstpos, axis = 1),
221
+ *[np.repeat(contribs[:, x], midpos, axis = 1) for x in [np.arange(P_length)[1:-1]]],
222
+ np.repeat(contribs[:, [-1]], lastpos, axis = 1)]
223
+
224
+ contribs_perpos = np.concatenate(contribs_perpos, axis = 1)
225
+
226
+ return contribs_perpos
227
+
228
+
229
+ def sppm_contrib(kppm_contribs_exp, kppm_idxs):
230
+
231
+ #This returns a list since they can be different sizes.
232
+
233
+ sizo = kppm_contribs_exp.shape[1]
234
+
235
+ outs = []
236
+ for idxs in kppm_idxs:
237
+
238
+ lex = len(idxs)
239
+ sex = sizo-lex
240
+
241
+ contrib_cuts = np.array([kppm_contribs_exp[do, ido: sex+ido]
242
+ for ido, do in enumerate(idxs)])
243
+
244
+ outs.append(np.sum(contrib_cuts, axis = 0))
245
+
246
+ return outs
247
+
248
+ def sppm_avgrevcomp(sppm_contribs):
249
+
250
+ # This averages the revcomps, flips the revcomp first tho
251
+ # sppm_contribs are of shape (contribs, len)
252
+ lex = len(sppm_contribs)
253
+ idxs = np.arange(lex)
254
+ sppm1 = [sppm_contribs[ido] for ido in idxs[:lex//2]]
255
+ sppm2 = [sppm_contribs[ido][::-1] for ido in idxs[lex//2:]]
256
+
257
+ sppm_avg = [(s1 + s2) / 2 for s1, s2 in zip(sppm1, sppm2)]
258
+
259
+ return sppm_avg
260
+
261
+ def contrib_fill(contribs, seq_len,
262
+ fill_with = np.nan, fill_center = True):
263
+
264
+ # contribs is an array or a list
265
+ expocons = []
266
+ for cont in contribs:
267
+ diffx = seq_len - len(cont)
268
+
269
+ diffx1 = diffx // 2 if fill_center else 0
270
+ diffx2 = diffx - diffx1
271
+
272
+ contexp = np.hstack([np.full(diffx1, fill_with), cont, np.full(diffx2, fill_with)])
273
+
274
+ expocons.append(contexp)
275
+
276
+ return np.array(expocons)
277
+
278
+
279
+
280
+ def kmer_cluster(inp, H_size,
281
+ num_clusters = 10, sample = None,
282
+ chunk = False,
283
+ cluster_mode = [KMeans, {}]):
284
+
285
+ sig_win_rs = tu.TridentWindow(inp, H_size)
286
+
287
+ num_sigs = sig_win_rs.shape[1]
288
+
289
+ lk = len(sig_win_rs)
290
+
291
+ if isinstance(sample, int):
292
+ rando = np.random.choice(np.arange(lk), sample, replace=False)
293
+ else: rando = np.arange(lk)
294
+
295
+ if chunk: #EACH SIGNAL INDIVIDUAL CLUSTER
296
+
297
+ clustod = []
298
+ for i in np.arange(num_sigs):
299
+
300
+ sig = sig_win_rs[:, i].reshape(lk, -1)
301
+ sig_samp = sig[rando]
302
+ clust = cluster_mode[0](n_clusters = num_clusters, **cluster_mode[1])
303
+ clust.fit(sig_samp)
304
+ clustod.append(clust.predict(sig))
305
+ print(f'finished: {i}')
306
+
307
+ clustod = np.stack(clustod, axis = 1)
308
+
309
+ else:
310
+ clust = cluster_mode[0](n_clusters = num_clusters, **cluster_mode[1])
311
+ clust.fit(sig_win_rs[rando].reshape(len(rando), -1))
312
+ clustod = clust.predict(sig_win_rs.reshape(lk, -1))
313
+
314
+ return clustod
315
+
316
+
317
+ def cluster_centroids(signal, clusters):
318
+ #signal is a (n, *) shape and clusters is a 1dim
319
+ return np.stack([np.mean(signal[clusters == b], axis = 0) for b in np.unique(clusters)])
320
+
321
+
322
+ def cluster_refine(signal, clusters, S = 1, return_newcents = False):
323
+
324
+ #uses agglomerative clsutering on centroids
325
+ # S is sensitivity for kneed.
326
+
327
+ clust1_uni = np.unique(clusters)
328
+
329
+ cents1 = cluster_centroids(signal, clusters)
330
+ model1 = AgglomerativeClustering(distance_threshold=0, n_clusters=None)
331
+ model1 = model1.fit(cents1.reshape(len(cents1), -1))
332
+
333
+ kn = KneeLocator(np.arange(len(model1.distances_)), model1.distances_,
334
+ S = S,
335
+ curve='convex', direction='increasing')
336
+ kno = kn.knee
337
+
338
+ model2 = AgglomerativeClustering(distance_threshold=model1.distances_[kno], n_clusters=None)
339
+ clust2 = model2.fit_predict(cents1.reshape(len(cents1), -1)) #you have a lst the length of unique clusters
340
+
341
+
342
+ clust2_uni = np.unique(clust2)
343
+
344
+ print(f'from {len(clust1_uni)} to {len(clust2_uni)}')
345
+
346
+ clust2_sepidx = [clust1_uni[clust2 == u] for u in clust2_uni] #This now tells you which of the original clusters belongs to each new cluster
347
+
348
+
349
+ newclust = np.ones(len(clusters))
350
+ for u,g in zip(clust2_uni, clust2_sepidx):
351
+ for h in g:
352
+ newclust[clusters == h] = u
353
+
354
+
355
+ newcents = cluster_centroids(signal, newclust)
356
+
357
+ return (newclust, newcents) if return_newcents else newclust
358
+
359
+
360
+ def centroid_contrib(wins, contribs, clusters,
361
+ weighted = True, score_mode = 'recipeuc',
362
+ chunk = False, return_centroids = False):
363
+ #wins shape = (number, signal, k, 1)
364
+ #contribs shape = (number, signal, pos, 1)
365
+ #clusters shape = (number, signal)
366
+
367
+ #this gets the centroids and their score per position
368
+
369
+ num_sigs = contribs.shape[1]
370
+
371
+ centroids_all = []
372
+ contrib_cents_all = []
373
+
374
+ if chunk:
375
+
376
+ for s in np.arange(num_sigs):
377
+
378
+ centroids = []
379
+ contrib_cents = []
380
+
381
+ clus = np.unique(clusters[:, s])
382
+
383
+ for c in clus:
384
+
385
+ wino = wins[:, s][clusters[:, s] == c] #this is now (num, k, 1)
386
+
387
+ centroid = np.mean(wino, axis = 0) #this is now (k, 1)
388
+ centroids.append(centroid)
389
+
390
+ contribo = contribs[:, s][clusters[:, s] == c] #this is now (num, pos, 1)
391
+
392
+ if weighted:
393
+
394
+ centroid_exp = np.expand_dims(centroid, axis = 0)
395
+
396
+ if score_mode == 'dot':
397
+ scores = np.sum(wino * centroid_exp, axis = (-1, -2))
398
+
399
+ elif score_mode == 'recipeuc':
400
+ scores = 1 / (np.sum(np.abs(wino-centroid_exp)**2, axis = (-1, -2))**(1/2))
401
+
402
+ scores = scores.reshape(-1, 1, 1)
403
+
404
+ contrib_cent = np.sum(contribo * scores, axis = 0) / np.sum(scores)
405
+
406
+ else:
407
+
408
+ contrib_cent = np.mean(contribo, axis = 0)
409
+
410
+ contrib_cents.append(contrib_cent)
411
+
412
+ centroids_all.append(np.stack(centroids))
413
+ contrib_cents_all.append(np.stack(contrib_cents))
414
+
415
+ else:
416
+
417
+ for c in np.unique(clusters):
418
+
419
+ wino = wins[clusters == c] #this is now (num, s, k, 1)
420
+
421
+ centroid = np.mean(wino, axis = 0) #this is now (s, k, 1)
422
+ centroids_all.append(centroid)
423
+
424
+ contribo = contribs[clusters == c, 0] #this is now (num, pos, 1)
425
+
426
+ if weighted: #Compare the wino to the centroid. Closer one is weighted more.
427
+
428
+ centroid_exp = np.expand_dims(centroid, axis = 0)
429
+
430
+ if score_mode == 'dot':
431
+ scores = np.sum(wino * centroid_exp, axis = (-1, -2, -3))
432
+
433
+ elif score_mode == 'recipeuc':
434
+ scores = 1 / (np.sum(np.abs(wino-centroid_exp)**2, axis = (-1, -2, -3))**(1/2))
435
+
436
+ scores = scores.reshape(-1, 1, 1)
437
+
438
+ contrib_cent = np.sum(contribo * scores, axis = 0) / np.sum(scores)
439
+
440
+ else:
441
+
442
+ contrib_cent = np.mean(contribo, axis = 0)
443
+
444
+ contrib_cents_all.append(contrib_cent)
445
+
446
+ centroids_all, contrib_cents_all = [np.array(x)
447
+ for x in [centroids_all, contrib_cents_all]]
448
+
449
+ return (centroids_all, contrib_cents_all) if return_centroids else contrib_cents_all
450
+
451
+
@@ -0,0 +1,309 @@
1
+ import numpy as np
2
+
3
+ import poseigen_seaside.basics as se
4
+ import poseigen_trident.utils as tu
5
+ import poseigen_trident.prongs as tp
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
+
11
+
12
+ class Oyster(nn.Module):
13
+
14
+ def __init__(self,
15
+ dim_i = (1,249,4), dim_f = (1,1,1),
16
+ incl_only = None,
17
+
18
+ DS = False, DS_wwp_func = nn.AvgPool2d, DS_act = False,
19
+
20
+
21
+ kE_k = 15, kE_cf_m = 100, kE_cf_ns = 1,
22
+ kE_ck_base = 4, kE_OneByOne = False,
23
+ kE_ck_grouped = False,
24
+
25
+ P_ck = None, P_center = False,
26
+
27
+ O_mods = None, O_mods_ns = 0.3,
28
+ O_cf_pu_m = None, O_cf_ns = 1,
29
+ O_ck_base = None, O_pool_k2s = 0.5,
30
+ O_pool_func = None, O_actb4pool = True,
31
+ O_dropout = 0, O_bias = False,
32
+
33
+ activations = nn.ReLU(), activation_f = None,
34
+ batchnorm = 'before'):
35
+
36
+ super(Oyster, self).__init__()
37
+
38
+
39
+ # 25.01.25 OCU Version A
40
+ # "Head" == now "kE" for k-mer Embedding
41
+ # "C" == now "P" for position
42
+ # "D" == now "O" for output
43
+ #kE_cf_m == the filter multiplier.
44
+ #P_center == whether to add padding to center things or not.
45
+
46
+ #25.03.28: added incl_only which uses only a subset of the data.
47
+
48
+
49
+ self.incl_only = None
50
+ if incl_only is not None:
51
+ if len(incl_only) == dim_i[0]: pass
52
+ else:
53
+ dim_i = (len(incl_only), dim_i[1], dim_i[2])
54
+ self.incl_only = incl_only
55
+
56
+ print(dim_i)
57
+
58
+
59
+
60
+
61
+ self.Reflect = tu.ReflectLayer(dims = [-1, -2]) if DS else nn.Identity()
62
+
63
+ kE_cf_f = kE_cf_m * dim_i[0]
64
+
65
+ if (dim_i[1] - kE_k + 1) < dim_f[1]: kE_k = dim_i[1] - dim_f[1] + 1
66
+
67
+ dim_i_ref = (dim_i[0], dim_i[1], dim_i[2]*2) if DS else dim_i
68
+ dim_kE_width = 2 if DS else 1
69
+ dim_kE_length = dim_i[1] - kE_k + 1
70
+ dim_kE = (kE_cf_f, dim_kE_length, dim_kE_width)
71
+
72
+ kE_Prong_args = {'dim_i': dim_i_ref, 'dim_f': dim_kE,
73
+ 'mods': 0, 'mods_ns': 0,
74
+ 'cf_i': None, 'cf_ns': kE_cf_ns, 'ck_base': kE_ck_base,
75
+ 'doublestranded': DS, 'OneByOne': kE_OneByOne,
76
+ 'ck_grouped': kE_ck_grouped,
77
+ 'activations': activations, 'activation_f': None,
78
+ 'batchnorm': batchnorm, 'dropout': None, 'bias': True,
79
+ 'out': False}
80
+
81
+ self.kE = nn.Sequential(*tp.Prong_X(**kE_Prong_args))
82
+
83
+ ################################################################
84
+
85
+ anti_ref_layers = [tu.FlipLayer(), tu.WWPLayer(func = DS_wwp_func)]
86
+ if DS_act: anti_ref_layers.append(activations)
87
+
88
+ self.AntiReflect = nn.Sequential(*anti_ref_layers) if DS else nn.Identity()
89
+
90
+ ###########################################################################################
91
+
92
+ P_func = nn.AvgPool2d
93
+
94
+ if P_ck == None: P_ck = 0
95
+ if P_ck == 0:
96
+ # Now this means do not pool.
97
+ self.P, P_length = (nn.Identity(), dim_kE_length)
98
+
99
+ elif P_ck > 1:
100
+
101
+ needed = P_ck - (dim_kE_length % P_ck)
102
+ pad2add = (needed // 2) if P_center else 0
103
+
104
+ self.P = P_func((P_ck, 1), stride = (P_ck,1),
105
+ padding = (pad2add, 0), ceil_mode = True, divisor_override = P_ck)
106
+
107
+ P_length = int(np.ceil(dim_kE_length / P_ck))
108
+
109
+ else: # means its 1, means we global pool
110
+
111
+ self.P = P_func((dim_kE_length, 1), stride = (1,1),
112
+ padding = (0, 0), ceil_mode = True, divisor_override = dim_kE_length)
113
+
114
+ P_length = 1
115
+
116
+ if O_cf_pu_m == None: O_cf_pu_m = 0
117
+ O_cf_pu = O_cf_pu_m * kE_cf_f if O_cf_pu_m != 0 else None
118
+
119
+ O_Prong_args = {'dim_i': (kE_cf_f, P_length, 1), 'dim_f': dim_f,
120
+ 'mods': O_mods, 'mods_ns': O_mods_ns,
121
+ 'cf_i': None, 'cf_pu': O_cf_pu,
122
+ 'cf_ns': O_cf_ns, 'ck_base': O_ck_base,
123
+ 'doublestranded': False,
124
+ 'ck_dynamic': False, 'skip_first_ck': False,
125
+ 'pool_k2s': O_pool_k2s, 'pool_func': O_pool_func, 'actb4pool': O_actb4pool,
126
+ 'activations': nn.ReLU(), 'activation_f': None,
127
+ 'batchnorm': batchnorm, 'dropout': O_dropout, 'bias': O_bias,
128
+ 'out': True}
129
+
130
+ self.O = nn.Sequential(*tp.Prong_Y(**O_Prong_args))
131
+
132
+ self.actf = nn.Identity() if activation_f is None else activation_f
133
+
134
+ def forward(self, x):
135
+
136
+ if self.incl_only is not None: x = x[:, self.incl_only]
137
+
138
+ x = self.Reflect(x)
139
+ x = self.kE(x)
140
+ x = self.AntiReflect(x)
141
+ x = self.P(x)
142
+ x = self.O(x)
143
+ x = self.actf(x)
144
+
145
+ return x
146
+
147
+
148
+
149
+ oys_args = {'dim_i': (1,249,4), 'dim_f': (1,1,1),
150
+ 'DS': False, 'DS_wwp_func': nn.AvgPool2d, 'DS_act': False,
151
+ 'incl_only': None,
152
+
153
+ 'kE_k': 15, 'kE_cf_m': 100, 'kE_cf_ns': 1,
154
+ 'kE_ck_base': 4, 'kE_OneByOne': False,
155
+ 'kE_ck_grouped': False,
156
+
157
+ 'P_ck': None, 'P_center': False,
158
+
159
+ 'O_mods': None, 'O_mods_ns': 0.3,
160
+ 'O_cf_pu_m': None, 'O_cf_ns': 1,
161
+ 'O_ck_base': None, 'O_pool_k2s': 0.5,
162
+ 'O_pool_func': None, 'O_actb4pool': True, 'O_dropout': 0,
163
+
164
+ 'activations': nn.ReLU(), 'activation_f': None,
165
+ 'batchnorm': 'before'}
166
+
167
+
168
+
169
+
170
+ class DualOyster(nn.Module):
171
+
172
+ def __init__(self,
173
+ dim_f = (1,1,1),
174
+
175
+ A_dim_i = (1,249,4), A_DS = False, A_incl_only = None,
176
+ A_kE_k = 15, A_kE_cf_m = 100, A_kE_cf_ns = 1,
177
+ A_kE_ck_base = 4, A_kE_OneByOne = False,
178
+ A_kE_ck_grouped = False,
179
+ A_P_ck = None,
180
+
181
+ use_B = True,
182
+
183
+ B_dim_i = (1,249,4), B_DS = False, B_incl_only = None,
184
+ B_kE_k = 15, B_kE_cf_m = 100, B_kE_cf_ns = 1,
185
+ B_kE_ck_base = 4, B_kE_OneByOne = False,
186
+ B_kE_ck_grouped = False,
187
+ B_P_ck = None,
188
+
189
+ DS_wwp_func = nn.AvgPool2d, DS_act = False,
190
+ P_center = False,
191
+ O_mods = None, O_mods_ns = 0.3,
192
+ O_cf_pu_m = None, O_cf_ns = 1,
193
+ O_ck_base = None, O_pool_k2s = 0.5,
194
+ O_pool_func = None, O_actb4pool = True, O_dropout = 0,
195
+ activations = nn.ReLU(), activation_f = None,
196
+ batchnorm = 'before'):
197
+
198
+ super(DualOyster, self).__init__()
199
+
200
+ oys_shared_args = {'dim_f': dim_f,
201
+ 'DS_wwp_func': DS_wwp_func, 'DS_act': DS_act,
202
+ 'P_center': P_center,
203
+ 'O_mods': O_mods, 'O_mods_ns': O_mods_ns,
204
+ 'O_cf_pu_m': O_cf_pu_m, 'O_cf_ns': O_cf_ns,
205
+ 'O_ck_base': O_ck_base, 'O_pool_k2s': O_pool_k2s,
206
+ 'O_pool_func': O_pool_func, 'O_actb4pool': O_actb4pool, 'O_dropout': O_dropout,
207
+ 'activations': activations,'batchnorm': batchnorm}
208
+
209
+
210
+ oys_A_args = {'dim_i': A_dim_i, 'DS': A_DS, 'incl_only': A_incl_only,
211
+ 'kE_k': A_kE_k, 'kE_cf_m': A_kE_cf_m, 'kE_cf_ns': A_kE_cf_ns,
212
+ 'kE_ck_base': A_kE_ck_base, 'kE_OneByOne': A_kE_OneByOne,
213
+ 'kE_ck_grouped': A_kE_ck_grouped,
214
+ 'P_ck': A_P_ck,
215
+ **oys_shared_args}
216
+
217
+ self.OysterA = Oyster(**oys_A_args)
218
+
219
+ self.use_B = use_B
220
+ if B_incl_only is not None:
221
+ if len(B_incl_only) == 0:
222
+ self.use_B = False
223
+ B_incl_only = None
224
+ if self.use_B == False: print('NOT USING!!!!')
225
+
226
+
227
+
228
+ oys_B_args = {'dim_i': B_dim_i, 'DS': B_DS, 'incl_only': B_incl_only,
229
+ 'kE_k': B_kE_k, 'kE_cf_m': B_kE_cf_m, 'kE_cf_ns': B_kE_cf_ns,
230
+ 'kE_ck_base': B_kE_ck_base, 'kE_OneByOne': B_kE_OneByOne,
231
+ 'kE_ck_grouped': B_kE_ck_grouped,
232
+ 'P_ck': B_P_ck,
233
+ **oys_shared_args}
234
+
235
+ self.OysterB = Oyster(**oys_B_args)
236
+
237
+ self.FA = activation_f if activation_f != None else nn.Identity()
238
+
239
+ def forward(self, x1, x2):
240
+
241
+ x1 = self.OysterA(x1)
242
+ x2 = self.OysterB(x2) if self.use_B else 0
243
+
244
+ return self.FA(x1 + x2)
245
+
246
+
247
+
248
+
249
+ DualOyster_BaseDict = {'A_dim_i': [[(1,249,4)],'cat'], #Set as a single dim_i
250
+ 'B_dim_i': [[(11,40,1)],'cat'], #Set as a single dim_i
251
+ 'dim_f': [[(2,1,1)],'cat'], #Set as a single dim_f
252
+
253
+ 'A_kE_ck_grouped': [[True, False], 'cat'], #NEED TO BE SET
254
+ 'A_DS': [[True, False], 'cat'],
255
+ 'A_kE_k': [[12, 20], 'int'],
256
+ 'A_kE_cf_m': [[50, 150], 'int'],
257
+ 'A_kE_cf_ns': [[1, 1.1, 1.2], 'cat'],
258
+ 'A_kE_ck_base': [[4, 5, 6, 7, 8, None], 'cat'],
259
+ 'A_kE_OneByOne': [[True, False], 'cat'],
260
+ 'A_P_ck': [[6, 20], 'int'],
261
+
262
+ 'B_kE_ck_grouped': [[True, False], 'cat'], #NEED TO BE SET
263
+ 'B_DS': [[True, False], 'cat'],
264
+ 'B_kE_k': [[12, 20], 'int'],
265
+ 'B_kE_cf_m': [[50, 150], 'int'],
266
+ 'B_kE_cf_ns': [[1, 1.1, 1.2], 'float'],
267
+ 'B_kE_ck_base': [[4, 5, 6, 7, 8, None], 'cat'],
268
+ 'B_kE_OneByOne': [[True, False], 'cat'],
269
+ 'B_P_ck': [[6, 20], 'int'],
270
+
271
+
272
+
273
+ 'DS_wwp_func': [[nn.MaxPool2d], 'cat'],
274
+ 'DS_act': [[False], 'cat'],
275
+
276
+ 'P_center': [[True, False], 'cat'],
277
+
278
+ 'O_mods': [[None], 'cat'],
279
+ 'O_mods_ns': [[0.5], 'cat'],
280
+ 'O_cf_pu_m': [[None], 'cat'],
281
+ 'O_cf_ns': [[1, 1.1, 1.2], 'cat'],
282
+ 'O_ck_base': [[None], 'cat'],
283
+ 'O_pool_k2s': [[0, 0.5], 'cat'],
284
+ 'O_pool_func': [[nn.MaxPool2d], 'cat'],
285
+ 'O_actb4pool': [[True], 'cat'],
286
+ 'O_dropout': [[0, 0.2], 'cat'],
287
+
288
+
289
+ 'activations': [[nn.ReLU(), nn.LeakyReLU()], 'cat'],
290
+ 'batchnorm': [[None, 'before'], 'cat']}
291
+
292
+
293
+ def Reset_DualOyster(dualoyster):
294
+ #ds consists of a "conv" and a "dense" module. Need to go through each one, see if its a conv and reset if so.
295
+
296
+
297
+ for oyster in [dualoyster.OysterA, dualoyster.OysterB]:
298
+ lke, los = len(oyster.kE), len(oyster.O)
299
+
300
+ for i in np.arange(lke):
301
+ if isinstance(oyster.kE[i], nn.Conv2d):
302
+ oyster.kE[i].reset_parameters()
303
+
304
+ for i in np.arange(los):
305
+ if isinstance(oyster.O[i], nn.Conv2d):
306
+ oyster.O[i].reset_parameters()
307
+ print('done reset mod')
308
+
309
+ return dualoyster