Interpretable learning of temporal cellular dynamics from single-cell data.

1. import torch
2. import torch.nn as nn
3. import torch.nn.functional as F
4. from torchdiffeq import odeint
5. from typing import Optional
6. from typing_extensions import Literal
8. from .module import LatentODE, Encoder, Decoder
11. class TNODE(nn.Module):
12. """
13. Class to automatically infer treatment wise cellular dynamics using autoencoder and neural ODE.
15. Parameters
16. ----------
17. n_int
18. Number of genes.
19. n_latent
20. Dimension of latent space.
21. (Default: 20)
22. n_sample
23. Number of samples or treatment as independant ODEs
24. (Default: 1)
25. n_ode_hidden
26. The dimensionality of the hidden layer for the ODE function.
27. (Default: 128)
28. n_vae_hidden
29. The dimensionality of the hidden layer for the AutoEncoder.
30. (Default: 128)
31. batch_norm
32. Whether to include `BatchNorm` layer.
33. (Default: 'False')
34. ode_method
35. ODE Solver method.
36. (Default: 'euler')
37. same_ode
38. Whether to initialize all ODEs with same initialization
39. (Default: 'False')
40. """
42. def __init__(
43. self,
44. n_int: int,
45. n_latent: int = 20,
46. n_sample: int = 1,
47. n_ode_hidden: int = 128,
48. n_vae_hidden: int = 128,
49. batch_norm: bool = False,
50. ode_method: str = 'dopri5',
51. same_ode: bool = False,
52. pre_ptime: bool = False,
53. reconstruct_xt: bool = False,
54. ):
55. super().__init__()
56. self.n_int = n_int
57. self.n_latent = n_latent
58. self.n_ode_hidden = n_ode_hidden
59. self.n_vae_hidden = n_vae_hidden
60. self.batch_norm = batch_norm
61. self.ode_method = ode_method
62. self.n_sample = n_sample
63. self.same_ode = same_ode
64. self.pre_ptime = pre_ptime
65. self.reconstruct_xt = reconstruct_xt
67. self.lode_func = nn.ModuleDict({f'node_{i}': LatentODE(n_latent, n_ode_hidden) for i in range(self.n_sample)})
68. self.encoder = Encoder(n_int, n_latent, n_vae_hidden, batch_norm)
69. self.decoder = Decoder(n_int, n_latent, n_vae_hidden, batch_norm)
70. self.beta, self.lam = torch.nn.Parameter(torch.Tensor(self.n_latent), requires_grad=True), torch.nn.Parameter(torch.Tensor(self.n_latent), requires_grad=True)
71. nn.init.uniform_(self.beta)
72. nn.init.uniform_(self.lam)
73. def forward(self, s: torch.Tensor, u: torch.Tensor, Pt: torch.Tensor, g: torch.Tensor) -> tuple:
74. """
75. Given the transcriptomes and the treatments of cells, this function predicts the time, latent space and dynamics of the cells.
77. Parameters
78. ----------
79. s
80. Spliced reads.
81. u
82. Unspliced reads.
83. Pt
84. Pseudotime tensor (Only if you want your velocity to follow specific pseudotime)
85. g
86. Sample/Treatment index.
88. Returns
89. ----------
90. 5-tuple of :class:`torch.Tensor`
91. Tensors for loss, including:
92. 1) total loss,
93. 2) reconstruction loss from encoder-derived latent space,
94. 3) reconstruction loss from ODE-solver latent space,
95. """
96. if self.pre_ptime:
97. Ts, z_s = self.encoder(s)
98. _, z_u = self.encoder(u)
99. Ts = Ts.ravel()
100. index = torch.argsort(Ts)
101. Ts = Ts[index]
102. s = s[index]
103. z_s = z_s[index]
104. u = u[index]
105. z_u = z_u[index]
106. g = g[index]
107. Pt = Pt[index]
109. index2 = (Ts[:-1] != Ts[1:])
110. index2 = torch.cat((index2, torch.tensor([True]).to(index2.device))) ## index2 is used to get unique time points as odeint requires strictly increasing/decreasing time points
111. Ts = Ts[index2]
112. s = s[index2]
113. z_s = z_s[index2]
114. z_u = z_u[index2]
115. u = u[index2]
116. g = g[index2]
117. Pt = Pt[index2]
118. pt_loss = F.mse_loss(Pt, Ts, reduction='mean')
119. else:
120. Ts, z_s = self.encoder(s)
121. _, z_u = self.encoder(u)
122. Ts = Ts.ravel()
123. index = torch.argsort(Ts)
124. Ts = Ts[index]
125. s = s[index]
126. z_s = z_s[index]
127. u = u[index]
128. z_u = z_u[index]
129. g = g[index]
130. index2 = (Ts[:-1] != Ts[1:])
131. index2 = torch.cat((index2, torch.tensor([True]).to(index2.device))) ## index2 is used to get unique time points as odeint requires strictly increasing/decreasing time points
132. Ts = Ts[index2]
133. s = s[index2]
134. z_s = z_s[index2]
136. z_u = z_u[index2]
137. u = u[index2]
138. g = g[index2]
139. pt_loss = 0
140. pred_z = torch.empty(s.size(0), self.n_latent).to(s.device)
141. z_div = torch.Tensor([0]).to(s.device)
142. for i in range(self.n_sample):
143. mask = g == i
145. if mask.sum() == 0:
146. continue
148. zsm = z_s[mask]
149. zum = z_u[mask]
150. Tm = Ts[mask]
151. z0 = zsm[0]
154. pred_z[mask] = odeint(self.lode_func[f'node_{i}'], z0, Tm, method = self.ode_method).view(-1, self.n_latent)
155. z_div += F.mse_loss(self.lode_func[f'node_{i}'](Tm, pred_z[mask]), (torch.exp(self.beta)*zum-torch.exp(self.lam)*zsm), reduction='mean')
156. pred_x_s = self.decoder(z_s)
157. pred_x_u = self.decoder(z_u)
158. if self.reconstruct_xt:
159. pred_x_t = self.decoder(pred_z)
160. recon_loss_xt = F.mse_loss(s, pred_x_t, reduction='mean')
161. else:
162. recon_loss_xt = F.mse_loss(z_s, pred_z, reduction='mean')
163. recon_loss_ec = F.mse_loss(s, pred_x_s, reduction='mean')
164. recon_loss_ec_u = F.mse_loss(u, pred_x_u, reduction='mean')
166. loss =recon_loss_ec + recon_loss_ec_u + z_div + recon_loss_xt + pt_loss
167. return loss, recon_loss_ec, recon_loss_ec_u, z_div

model.py at commit 0f069ba, under BSD-3-Clause · at the source