Skip to content

Commit 3dcc954

Browse files
trishortstrishorts
authored
Merge branch 'master' into pride-download
2 parents 9249f7e + 00bc5cf commit 3dcc954

84 files changed

Lines changed: 8613 additions & 277 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎mzLib/PredictionClients/Koina/AbstractClasses/CollisionalCrossSectionModel.cs‎

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
using Omics.SequenceConversion;
2+
using PredictionClients.Koina.Client;
3+
using PredictionClients.Koina.Interfaces;
4+
using PredictionClients.Koina.Util;
5+
using System.ComponentModel;
6+
using MzLibUtil;
7+
using Easy.Common.Extensions;
8+
9+
namespace PredictionClients.Koina.AbstractClasses
10+
{
11+
/// <summary>
12+
/// Represents a collisional cross section prediction result for a single peptide.
13+
/// </summary>
14+
public record PeptideCCSPrediction(
15+
string FullSequence,
16+
string ValidatedFullSequence,
17+
int PrecursorCharge,
18+
double? PredictedCCS,
19+
WarningException? Warning = null
20+
);
21+
22+
/// <summary>
23+
/// Input parameters for CCS prediction models.
24+
/// </summary>
25+
public record CCSPredictionInput(
26+
string FullSequence,
27+
int PrecursorCharge
28+
)
29+
{
30+
public string? ValidatedFullSequence { get; set; }
31+
public WarningException? SequenceWarning { get; set; }
32+
public WarningException? ParameterWarning { get; set; }
33+
}
34+
35+
/// <summary>
36+
/// Abstract base class for collisional cross section (CCS) prediction models using the Koina API.
37+
///
38+
/// Thread safety: instances are NOT thread-safe. Predict and related methods
39+
/// mutate instance state (ModelInputs, ValidInputsMask, Predictions); callers must not invoke
40+
/// these methods concurrently on the same instance, nor read Predictions while a call is in flight.
41+
/// Use one instance per concurrent caller (or serialize externally) when sharing across pipelines.
42+
/// </summary>
43+
public abstract class CollisionalCrossSectionModel : KoinaModelBase<CCSPredictionInput, PeptideCCSPrediction>, IPredictor<CCSPredictionInput, PeptideCCSPrediction>
44+
{
45+
protected CollisionalCrossSectionModel(ISequenceConverter sequenceConverter)
46+
: base(sequenceConverter)
47+
{
48+
}
49+
50+
public virtual HashSet<int>? AllowedPrecursorCharges => new() { 1, 2, 3, 4, 5, 6 };
51+
public override IReadOnlySet<int> AllowedUnimodIds => new HashSet<int>();
52+
public override abstract SequenceConversionHandlingMode ModHandlingMode { get; init; }
53+
public virtual IncompatibleParameterHandlingMode ParameterHandlingMode { get; init; }
54+
55+
public List<CCSPredictionInput> ModelInputs { get; protected set; } = new();
56+
public bool[] ValidInputsMask { get; protected set; } = Array.Empty<bool>();
57+
public List<PeptideCCSPrediction> Predictions { get; protected set; } = new();
58+
59+
protected virtual async Task<List<PeptideCCSPrediction>> AsyncThrottledPredictor(List<CCSPredictionInput> modelInputs)
60+
{
61+
if (modelInputs.IsNullOrEmpty())
62+
{
63+
Predictions = new List<PeptideCCSPrediction>();
64+
return Predictions;
65+
}
66+
67+
ModelInputs = modelInputs;
68+
ValidInputsMask = new bool[ModelInputs.Count];
69+
var validInputs = new List<CCSPredictionInput>();
70+
71+
for (int i = 0; i < ModelInputs.Count; i++)
72+
{
73+
var cleanedSequence = TryCleanSequence(ModelInputs[i].FullSequence, out var apiSequence, out var modHandlingWarning);
74+
var validModelParams = ValidateModelSpecificInputs(ModelInputs[i], out var parameterWarning);
75+
if (cleanedSequence != null && apiSequence != null && validModelParams)
76+
{
77+
ModelInputs[i] = ModelInputs[i] with { ValidatedFullSequence = apiSequence, SequenceWarning = modHandlingWarning, ParameterWarning = parameterWarning };
78+
ValidInputsMask[i] = true;
79+
validInputs.Add(ModelInputs[i]);
80+
}
81+
else
82+
{
83+
ModelInputs[i] = ModelInputs[i] with { ValidatedFullSequence = null, SequenceWarning = modHandlingWarning, ParameterWarning = parameterWarning };
84+
ValidInputsMask[i] = false;
85+
}
86+
}
87+
88+
var predictions = new List<PeptideCCSPrediction>();
89+
if (validInputs.Count > 0)
90+
{
91+
var batchedRequests = ToBatchedRequests(validInputs);
92+
var batchChunks = batchedRequests.Chunk(MaxNumberOfBatchesPerRequest).ToList();
93+
int sessionTimeoutInMinutes = (int)Math.Ceiling((batchedRequests.Count * 2 * BenchmarkedTimeForOneMaxBatchSizeInMilliseconds + ThrottlingDelayInMilliseconds * batchChunks.Count) / 6e4);
94+
sessionTimeoutInMinutes = Math.Max(sessionTimeoutInMinutes, 1);
95+
96+
var responses = new List<string>();
97+
using var cts = new CancellationTokenSource(TimeSpan.FromMinutes(sessionTimeoutInMinutes));
98+
99+
for (int i = 0; i < batchChunks.Count; i++)
100+
{
101+
var batchChunk = batchChunks[i];
102+
var responseChunk = await Task.WhenAll(batchChunk.Select(request => SendInferenceRequestAsync(ModelName, request, cts.Token)));
103+
responses.AddRange(responseChunk);
104+
105+
if (i < batchChunks.Count - 1)
106+
{
107+
await Task.Delay(ThrottlingDelayInMilliseconds);
108+
}
109+
}
110+
111+
predictions = ResponseToPredictions(responses, validInputs);
112+
}
113+
114+
var realignedPredictions = new List<PeptideCCSPrediction>();
115+
int predictionIndex = 0;
116+
for (int i = 0; i < ValidInputsMask.Length; i++)
117+
{
118+
if (ValidInputsMask[i])
119+
{
120+
realignedPredictions.Add(predictions[predictionIndex]);
121+
predictionIndex++;
122+
}
123+
else
124+
{
125+
realignedPredictions.Add(new PeptideCCSPrediction(
126+
FullSequence: ModelInputs[i].FullSequence,
127+
ValidatedFullSequence: null,
128+
PrecursorCharge: ModelInputs[i].PrecursorCharge,
129+
PredictedCCS: null,
130+
Warning: ModelInputs[i].ParameterWarning ?? ModelInputs[i].SequenceWarning ?? new WarningException("Input was invalid and skipped during prediction.")
131+
));
132+
}
133+
}
134+
135+
Predictions = realignedPredictions;
136+
return Predictions;
137+
}
138+
139+
private readonly object _predictLock = new();
140+
141+
public List<PeptideCCSPrediction> Predict(List<CCSPredictionInput> modelInputs)
142+
{
143+
// Offload to Task.Run so the awaited continuations inside AsyncThrottledPredictor run on
144+
// the ThreadPool (no SynchronizationContext) rather than trying to resume on a blocked
145+
// caller thread. Calling .GetAwaiter().GetResult() directly would deadlock under a
146+
// single-threaded SynchronizationContext (WinForms/WPF/ASP.NET non-Core). The lock
147+
// serializes concurrent callers against the shared instance state. See
148+
// FragmentIntensityModel.Predict for the full rationale (stopgap Option A).
149+
lock (_predictLock)
150+
{
151+
return Task.Run(() => AsyncThrottledPredictor(modelInputs)).GetAwaiter().GetResult();
152+
}
153+
}
154+
155+
protected virtual bool ValidateModelSpecificInputs(CCSPredictionInput input, out WarningException? warning)
156+
{
157+
warning = null;
158+
159+
// Always reject non-physical charges (< 1) regardless of AllowedPrecursorCharges.
160+
if (input.PrecursorCharge < 1)
161+
{
162+
string message = $"Precursor charge {input.PrecursorCharge} is not a valid physical charge. Charge must be positive.";
163+
switch (ParameterHandlingMode)
164+
{
165+
case IncompatibleParameterHandlingMode.ThrowException:
166+
throw new ArgumentException(message);
167+
case IncompatibleParameterHandlingMode.ReturnNull:
168+
warning = new WarningException(message);
169+
return false;
170+
default:
171+
throw new ArgumentException($"Unhandled ParameterHandlingMode: {ParameterHandlingMode}");
172+
}
173+
}
174+
175+
// null = not applicable (skip), empty = any charge, populated = restrict to listed.
176+
if (AllowedPrecursorCharges != null && !AllowedPrecursorCharges.IsNullOrEmpty() && !AllowedPrecursorCharges.Contains(input.PrecursorCharge))
177+
{
178+
string exceptionMessage = $"Precursor charge {input.PrecursorCharge} is not supported by this model. Allowed precursor charges: {string.Join(", ", AllowedPrecursorCharges)}.";
179+
switch (ParameterHandlingMode)
180+
{
181+
case IncompatibleParameterHandlingMode.ThrowException:
182+
throw new ArgumentException(exceptionMessage);
183+
184+
case IncompatibleParameterHandlingMode.ReturnNull:
185+
warning = new WarningException(exceptionMessage);
186+
return false;
187+
default:
188+
throw new ArgumentException($"Unhandled ParameterHandlingMode: {ParameterHandlingMode}");
189+
}
190+
}
191+
192+
return true;
193+
}
194+
195+
protected virtual List<PeptideCCSPrediction> ResponseToPredictions(
196+
IReadOnlyList<string> responses,
197+
List<CCSPredictionInput> requestInputs)
198+
{
199+
var predictions = new List<PeptideCCSPrediction>();
200+
if (requestInputs.IsNullOrEmpty())
201+
{
202+
return predictions;
203+
}
204+
205+
var deserializedResponses = responses.Select(r => Newtonsoft.Json.JsonConvert.DeserializeObject<ResponseJSONStruct>(r)).ToList();
206+
if (deserializedResponses.IsNullOrEmpty() || deserializedResponses.Any(r => r == null))
207+
{
208+
throw new Exception("Something went wrong during deserialization of responses.");
209+
}
210+
211+
var ccsOutputs = deserializedResponses.SelectMany(r =>
212+
{
213+
var output = r!.Outputs.FirstOrDefault(o => o.Name == "ccs")
214+
?? throw new Exception($"API response is missing expected output 'ccs'. Found: {string.Join(", ", r.Outputs.Select(o => o.Name))}.");
215+
return output.Data.Select(d => Convert.ToDouble(d));
216+
}).ToList();
217+
218+
if (ccsOutputs.Count != requestInputs.Count)
219+
{
220+
throw new Exception("The number of CCS predictions does not match the number of input peptides.");
221+
}
222+
223+
for (int i = 0; i < requestInputs.Count; i++)
224+
{
225+
predictions.Add(new PeptideCCSPrediction(
226+
FullSequence: requestInputs[i].FullSequence,
227+
ValidatedFullSequence: requestInputs[i].ValidatedFullSequence,
228+
PrecursorCharge: requestInputs[i].PrecursorCharge,
229+
PredictedCCS: ccsOutputs[i],
230+
Warning: requestInputs[i].SequenceWarning
231+
));
232+
}
233+
234+
return predictions;
235+
}
236+
}
237+
}

0 commit comments

Comments
 (0)